pythonpython-typingpython-decorators

How to add type annotation to self parameter of the decorator of class method?


So, I have a decorator for a class like this

def something_todo(f):
    @functools.wraps(f)
    def decorator(self, *args, **kwargs):
        # do something
        return f(self, *args, **kwargs)
    return decorator

I want to have type annotation to self parameter in the decorator. But, writing like this is not working

def something_todo(f):
    @functools.wraps(f)
    def decorator(self: SomeClass, *args, **kwargs):
        # do something
        return f(self, *args, **kwargs)
    return decorator

class SomeClass:

    def __init__(self):
        # init here

    @something_todo
    def some_method(self, *args, **kwargs):
        # some process

And it even gets worse if I write the decorator and the class in different scripts. This will likely be happened

ImportError: cannot import name 'SomeClass' from partially initialized module 'someclass' (most likely due to a circular import)

The reason why I did this is for clarity so people know that my decorator is only for that class method. Also, I would easily check all methods or properties of the class in the editor without the need to open the script that contains the class.

EDIT

So this is the recap of the solution:

Let's say I have 2 scripts. One contains the class, and one contains the decorator function. For example,

# inside decorators.py
def something_todo(f):
    @functools.wraps(f)
    def decorator(self, *args, **kwargs):
        # do something
        return f(self, *args, **kwargs)
    return decorator

and

# inside someclass.py

from decorators import something_todo

class SomeClass:

    def __init__(self):
        # init here

    @something_todo
    def some_method(self, *args, **kwargs):
        # some process

If I want to add a type annotation to self parameter in the decorator, I couldn't just do this

# inside decorators.py

from someclass import SomeClass

def something_todo(f):
    @functools.wraps(f)
    def decorator(self: SomeClass, *args, **kwargs):
        # do something
        return f(self, *args, **kwargs)
    return decorator

Because it will raise ImportError when I tried to import the class inside another script. So, to prevent the error, I could just do this

# inside decorators.py

from someclass import *

def something_todo(f):
    @functools.wraps(f)
    def decorator(self: "SomeClass", *args, **kwargs):
        # do something
        return f(self, *args, **kwargs)
    return decorator

And it works fine.


Solution

  • Two useful tricks:

    1. Forward-declare types as string literals, e.g. self: 'SomeClass'
    2. Use callable protocols for more flexibility in defining callable TypeVars.
    import functools
    from typing import cast, Any, Callable, Protocol, TypeVar
    
    
    class SomeMethod(Protocol):
        def __call__(
            _self,
            self: 'SomeClass',
            *args: Any,
            **kwargs: Any
        ) -> Any: ...
    
    
    _SomeMethod = TypeVar("_SomeMethod", bound=SomeMethod)
    
    
    def something_todo(f: _SomeMethod) -> _SomeMethod:
        @functools.wraps(f)
        def wrapper(self: SomeClass, *args, **kwargs):
            # do something
            return f(self, *args, **kwargs)
        return cast(_SomeMethod, wrapper)
    
    
    class SomeClass:
    
        def __init__(self):
            # init here
            pass
    
        @something_todo
        def some_method(self, *args, **kwargs):
            # some process
            pass
    
    
    @something_todo
    def foo():
        print('error: Value of type variable "_SomeMethod" of "something_todo" cannot be "Callable[[], None]"')