Wrapping a decorator and preserving typing

Python Help
1

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

@retry(retry=retry_if_exception_type(duckdb.IOException),
       stop=stop_after_attempt(10),
       wait=wait_fixed(1),
       reraise=True)
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:

@retry_io
def foo():
    ...

This works at run time:

def retry_io():
    def decorator(func: WrappedFn) -> WrappedFn:
        @retry(retry=retry_if_exception_type(IOException),
               stop=stop_after_attempt(8),
               wait=wait_fixed(0.25),
               reraise=True)
        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…

2

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),
               stop=stop_after_attempt(8),
               wait=wait_fixed(0.25),
               reraise=True)(fn)
1 Like
3

Ah, thank you, that’s it!

I tweaked the signature:

def retry_io(fn: WrappedFn) -> WrappedFn:

where WrappedFn.

4

Even simpler, since functions are first-class objects:

retry_io = retry(
    retry=retry_if_exception_type(IOException),
    stop=stop_after_attempt(8),
    wait=wait_fixed(0.25),
    reraise=True
)
4 Likes
5

Love it! Feels like Haskell :slight_smile:

6

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:

 @retry
 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(
     retry,
     retry=retry_if_exception_type(IOException),
     ..........
 )

to curry the function.

7

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).