如何让Python IDE的类型检查识别类装饰器动态添加的方法

我最近遇到个头疼的问题——想用类装饰器给目标类动态混入新方法，同时还得让IDE的类型检查能同时识别原类和混入方法的类型提示。先给你看看我最开始的实现代码：

from typing import Type, TypeVar

T = TypeVar("T")


class NewMethodsMixin:
    def new_method(self) -> str:
        """This is a new method added by a mixin."""
        print("added by mixin")


def class_decorator(cls: Type[T]):
    class Decorated(NewMethodsMixin, cls):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            cls.__init__(self, *args, **kwargs)

    return Decorated


@class_decorator
class MyClass:
    def old_method(self) -> str:
        print("original method")


my_instance = MyClass()
my_instance.old_method()  # 'original method'
my_instance.new_method()  # 'added by mixin'

代码跑起来功能正常，但IDE里old_method完全没有类型提示，就好像IDE根本不知道这个实例原本属于MyClass一样。

后来我试着给装饰器加上返回类型标注-> Type[T]，改成这样：

from typing import Type, TypeVar

T = TypeVar("T")


class NewMethodsMixin:
    def new_method(self) -> str:
        """This is a new method added by a mixin."""
        print("added by mixin")


def class_decorator(cls: Type[T]) -> Type[T]:
    class Decorated(NewMethodsMixin, cls):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            cls.__init__(self, *args, **kwargs)

    return Decorated


@class_decorator
class MyClass:
    def old_method(self) -> str:
        print("original method")

这下好了，old_method的类型提示终于回来了，但新混入的new_method又没了踪影，IDE完全识别不到这个方法的存在。

我又试过用Union来标注返回类型：-> Type[T] | Type[NewMethodsMixin]，结果IDE倒是能识别两个方法的提示了，但如果我用装饰后的类去做继承的话，类型检查直接报错，完全没法正常用。

因为不想让用户去处理多继承的各种坑，我还是更倾向用装饰器的方式，折腾了半天终于找到靠谱的解决方案，分享给你：

方案一：使用交集类型（Intersection）标注装饰器返回值 #

从Python 3.10开始，官方支持了交集类型（&），可以用来表示一个类型同时属于多个类型的集合，配合TypeVar就能完美解决这个问题。如果你的项目还在使用Python 3.9及以下版本，可以用typing_extensions.Intersection来替代原生的交集语法。

修改后的代码如下：

from typing import Type, TypeVar
# Python 3.9及以下版本请替换成这行：
# from typing_extensions import Intersection

T = TypeVar("T")


class NewMethodsMixin:
    def new_method(self) -> str:
        """This is a new method added by a mixin."""
        print("added by mixin")


def class_decorator(cls: Type[T]) -> Type[T & NewMethodsMixin]:
    class Decorated(NewMethodsMixin, cls):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            # 注意：这里不需要手动调用cls.__init__，super()已经会按MRO顺序调用所有父类的初始化方法，删掉这行避免重复初始化
            # cls.__init__(self, *args, **kwargs)

    return Decorated


@class_decorator
class MyClass:
    def old_method(self) -> str:
        print("original method")


my_instance = MyClass()
my_instance.old_method()  # 现在有完整的类型提示
my_instance.new_method()  # 现在也有完整的类型提示

这样修改后，类型检查器会明确知道：装饰器返回的类同时是T（被装饰的原类）和NewMethodsMixin的子类，所以会同时识别两个类的所有方法，而且后续用这个类做继承也不会有类型问题。

方案二：使用Protocol定义合并后的方法集合（兼容旧Python版本） #

如果你的项目需要兼容Python 3.9及以下，或者不想用交集类型，还可以用Protocol来定义一个包含所有需要方法的协议，然后让装饰器返回这个协议的类型：

from typing import Type, TypeVar, Protocol

T = TypeVar("T")


class NewMethodsMixin:
    def new_method(self) -> str:
        """This is a new method added by a mixin."""
        print("added by mixin")


# 定义一个协议，包含原类和混入类的所有方法
class CombinedProtocol(Protocol):
    def old_method(self) -> str: ...
    def new_method(self) -> str: ...


def class_decorator(cls: Type[T]) -> Type[CombinedProtocol]:
    class Decorated(NewMethodsMixin, cls):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)

    return Decorated


@class_decorator
class MyClass:
    def old_method(self) -> str:
        print("original method")

不过这个方案有个小缺点：如果后续MyClass或者NewMethodsMixin的方法有变动，你需要手动同步更新CombinedProtocol里的方法定义，灵活性不如交集类型。

备注：内容来源于stack exchange，提问作者owl sandy