Python - change annotation in decorator

Question

Say I have a function that has type annotations:

def f(n: float) -> int:
    if random.random() >= n:
        raise Exception
    return 1

I want to decorate this function, and the decorated version will extend the return values of the original function:

def decorator(f):
    def decorated(max_tries: int, delay: int, *a, **kwa):  # annotations?
        for i in range(1, max_tries+1):
            try:
                return True, f(*a, **kwa)
            except Exception:
                log('try %d of %d failed' % (i, max_tries))
                time.sleep(delay)
        
        return False, None
    
    return decorated

In this example I want to change the original return type int to (bool, Union[int, None]).

I tried this:

def decorator(f):
    @functools.wraps(f)
    def decorated(max_tries: int, delay: int, *a, **kwa):
        for i in range(1, max_tries+1):
            try:
                return True, f(*a, **kwa)
            except Exception:
                log('failed try %d of %d' % (i, max_tries))
                time.sleep(delay)
        
        return False, None
    
    ret = f.__annotations__['return']
    decorated.__annotations__['return'] = (bool, Union([ret, None]))
    
    return decorated

And I get TypeError: Cannot instantiate <class 'typing.UnionMeta'>.

Does anyone know how to do this?

**EDIT:**

>>> type(ret)
<class 'typing.UnionMeta'>
>>> isinstance(ret, Union)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "C:\Program Files\Python35-32\lib\typing.py", line 563, in __instancecheck__
    raise TypeError("Unions cannot be used with isinstance().")
TypeError: Unions cannot be used with isinstance().

What exactly are classes from the typing module (eg. List, Union, etc.)? Are they instantiated when defining a function? Why cannot they be instantiated by hand?

**EDIT 2:**

Apparently the way to "instantiate" a Union is to use the same notation as in type annotations: Union[int, None].

def decorator(f):
    @functools.wraps(f)
    def decorated(max_tries: int, delay: int, *a, **kwa):
        for i in range(1, max_tries+1):
            try:
                return True, f(*a, **kwa)
            except Exception:
                log('failed try %d of %d' % (i, max_tries))
                time.sleep(delay)
        
        return False, None
    
    ret = f.__annotations__['return']
    decorated.__annotations__['return'] = (bool, Union[ret, None])
    
    return decorated

Now the problem is to add parameters max_tries: int and delay: int to the decorated function's annotations, and for these to be present in help(decorated_f). Any help with that?

Answer 1

You can do this by inspecting the function’s and/or messing with the wrapper’s __signature__, see inspect.signature(). If you just want the return parameter:

def decorator(f):
    orig = inspect.signature(f)

    def decorated(max_tries: int, delay: int,
                  *a, **kwa) -> tuple[bool, orig.return_annotation | None]:
        for i in range(1, max_tries+1):
            try:
                return True, f(*a, **kwa)
            except Exception:
                log('try %d of %d failed' % (i, max_tries))
                time.sleep(delay)

        return False, None

    return decorated

@decorator
def f(n: float) -> int:
    if random.random() >= n:
        raise Exception
    return 1

This gives a semi-satisfying answer as help(f):

decorated(max_tries: int, delay: int, *a, **kwa) -> tuple[bool, int | None]

With a couple more tweaks such as using functools.update_wrapper and replacing the *a, **kwa with the original parameters:

import inspect
import functools

def decorator(f):
    def decorated(max_tries: int, delay: int, *a, **kwa):
        for i in range(1, max_tries+1):
            try:
                return True, f(*a, **kwa)
            except Exception:
                log('try %d of %d failed' % (i, max_tries))
                time.sleep(delay)

        return False, None

    orig = inspect.signature(f)
    wrap = inspect.signature(decorated)

    functools.update_wrapper(decorated, f)
    decorated.__signature__ = orig.replace(
        parameters=list(wrap.parameters.values())[:-2] + list(orig.parameters.values()),
        return_annotation=tuple[bool, orig.return_annotation | None]
    )

    return decorated

@decorator
def f(n: float) -> int:
    """ Even has a docstring """
    if random.random() >= n:
        raise Exception
    return 1

We get from help(f):

f(max_tries: int, delay: int, n: float) -> tuple[bool, int | None]
    Even has a docstring