pythonpython-typingpython-decorators

How do I type hint for class wrapper or class with method assigned elsewhere than definition in Python3?


In some cases, it is required to inject the same method to multiple classes. To avoid repeated code, I can use a class decorator or class wrapper to assign new attributes and methods to specific class. (Since type hint for wrapper and decorator are the same, I will take decorator as example)

class Base:
    pass

# decorator
def decorator(cls: Input_Type) -> Output_Type:
    def new_method(self):
        pass
    cls.new_method = new_method
    return cls

# wrapper
def wrapper(cls: Input_Type) -> Output_Type:
    class WrappedClass(cls):
        def new_method(self):
            pass
    return WrappedClass

Moreover, for convenience and better maintainability, I want to add type hint for this decorator/wrapper (Output_Type in code block). By type hint, IDE should be able to achieve static analysis and hinting for both original class attr and newly injected attr. How should I do the type hint?

Versions

python==3.8.10
VScode==1.91.1
# VScode Extensions
Pylance==v2024.7.1
Python-Type-Hint==v1.5.1

Return original type

I've tried typing.TypeVar, like

import typing as t
T = t.TypeVar("T", bound=t.Type)

def func(cls: T) -> T:
    ...

However, in this way, the return type is exactly the same as input type, and newly assigned methods won't be hinted.

Using Inheritance

I've also tried using inheritance from typing.Generic and defining a Protocol class, like

import typing as t
T = t.TypeVar("T", bound=t.Type)

class Injected(t.Generic[T]):
    def new_method(self): ...

def func(cls: T) -> Injected[T]:
    ...

But in this way, IDE doesn't hint anything, neither new method nor the originals.

I wonder if there is a way to attach new methods to an existing type(class) in type hint.


Solution

  • You don't need a decorator to assign an attribute to a class, or to create a wrapper class. The simplest way to do what you need is simply to assign the method in the class scope. This also has the advantage of working well with static analysis tools.

    def method_implementation(self, arg: int) -> set[int]:
        """
        a great docstring
        """
        return set()
    
    class Foo:
        method = method_implementation
    
    class Bar:
        method = method_implementation
        def another_method(self, x:str) -> None:
            print(x)
    
    
    bar = Bar()
    bar.method("foo")
    

    And see, static analysis tools and IDEs will understand it. Here is mypy catching the correct error:

    (py312) Juans-MBP:test juan$ mypy foo.py 
    foo.py:17: error: Argument 1 has incompatible type "str"; expected "int"  [arg-type]
    Found 1 error in 1 file (checked 1 source file)
    

    And IDE intellisense will understand it, at least using pylance with VSCode:

    enter image description here

    enter image description here