Wrapping a decorator and preserving typing

I use a decorator from the tenacity library in lots of places, like this:

def foo():

I’d love to ‘rewrap’ that in a new decorator (i.e. that exact call, with those exact args), in e.g. @retry_io, so I can:

def foo():

This works at run time:

def retry_io():
    def decorator(func: WrappedFn) -> WrappedFn:
        def wrapper(*args: Any, **kwargs: Any):
            return func(*args, **kwargs)
        return wrapper
    return decorator

But the return wrapper line fails to type check with:

Expression of type "(*args: Any, **kwargs: Any) -> Any" cannot be assigned to return type "WrappedFn@decorator"
  Type "(*args: Any, **kwargs: Any) -> Any" cannot be assigned to type "WrappedFn@decorator"PylancereportReturnType

Any ideas, or a more elegant way to do this?
I feel like there’s some functools magic, but I don’t know it…

Remember that decorators are just functions like any other. So, you could do it like that:

def retry_io(fn):
    return retry(retry=retry_if_exception_type(IOException),
1 Like

Ah, thank you, that’s it!

I tweaked the signature:

def retry_io(fn: WrappedFn) -> WrappedFn:

where WrappedFn.

Even simpler, since functions are first-class objects:

retry_io = retry(

Love it! Feels like Haskell :slight_smile:

This is indeed the nice way to write it. But it isn’t because functions
are first class objects.

But in this instance this only works because retry() is a decorator
which accepts optional parameters… With no function position 0 argument
it returns a curried retry decorator. I’m assuming you can use in in a
default mode:

 def f(...):

although having glanced at the tenacity source code I’m less sure.

If we were treating retry just like a function you’d need:

 from functools import partial

 retry_io = partial(

to curry the function.

It was a lazy justification on my part, sure. What I did was basically eta-reduction, but in a way that also evaluates one step eagerly.

The approach shown would work just fine if the parameters were all mandatory. Decorators with parameters are basically pre-curried: calling them with the configuration arguments produces another function, which then acts like a parameter-less decorator (called with the function to return a transformed function). Felix’s approach is to wrap that double-call process and specify the arguments for the first call, while using a supplied value (i.e., the function to decorate) for the second call; this effectively transforms the decorator with parameters into one without. I simplified this by just making the first call now. Nothing about that relies on any secondary functionality that the retry decorator might have to allow parameter-less use (although it certainly is possible to create decorators that work this way using optional arguments).