类装饰器的类型提示

Question

我有一个类装饰器，它删除一个方法并将另一个方法添加到类中。

我如何提供类型提示？我显然已经尝试自己研究这个问题，但无济于事。

大多数人声称这需要交叉类型。有推荐的解决方案吗？我缺少什么吗？

示例代码：

class MyProtocol(Protocol):
    def do_check(self) -> bool:
        raise NotImplementedError

def decorator(clazz: type[MyProtocol]) -> ???:
    do_check: Callable[[MyProtocol], bool] = getattr(clazz, "do_check")

    def do_assert(self: MyProtocol) -> None:
        assert do_check(self)

    delattr(clazz, "do_check")
    setattr(clazz, "do_assert", do_assert)
    
    return clazz

@decorator
class MyClass(MyProtocol):
    def do_check(self) -> bool:
        return False

mc = MyClass()
mc.do_check() # hints as if exists, but doesn't
mc.do_assert() # no hints, but works

我想我正在寻找的是

decorator

的正确返回类型。

Answer 1

没有类型注释可以做你想做的事。即使使用交集类型，也没有办法表达删除属性的操作 - 您能做的最好的事情就是与用某种不可用的

descriptor

覆盖 do_check 的类型进行交集。

您所要求的可以通过 mypy 插件来完成。基本实现后，结果可能如下所示：

from package.decorator_module import MyProtocol, decorator

@decorator
class MyClass(MyProtocol):
    def do_check(self) -> bool:
        return False

>>> mc = MyClass()  # mypy: Cannot instantiate abstract class "MyClass" with abstract attribute "do_check" [abstract]
>>> mc.do_check()   # raises `NotImplementedError` at runtime
>>> mc.do_assert()  # OK

请注意，

mc.do_check

存在，但被插件检测为抽象方法。这与运行时实现相匹配，因为

delattr

删除

MyClass.do_check

只是公开了父级

MyProtocol.do_check

，而

typing.Protocol

上的非重写方法是抽象方法，您无法在不重写它们的情况下实例化该类。

---

这是一个基本的实现。使用以下目录结构：

project/
  mypy.ini
  mypy_plugin.py
  test.py
  package/
    __init__.py
    decorator_module.py

mypy.ini

的内容

[mypy]
plugins = mypy_plugin.py

mypy_plugin.py

的内容

from __future__ import annotations

import typing_extensions as t

import mypy.plugin
import mypy.plugins.common
import mypy.types

if t.TYPE_CHECKING:
    import collections.abc as cx
    import mypy.nodes

def plugin(version: str) -> type[DecoratorPlugin]:
    return DecoratorPlugin

class DecoratorPlugin(mypy.plugin.Plugin):

    # See https://mypy.readthedocs.io/en/stable/extending_mypy.html#current-list-of-plugin-hooks
    # Since this is a class definition modification with a class decorator
    # and the class body should have been semantically analysed by the time
    # the class definition is to be manipulated, we choose
    # `get_class_decorator_hook_2`
    def get_class_decorator_hook_2(
        self, fullname: str
    ) -> cx.Callable[[mypy.plugin.ClassDefContext], bool] | None:
        if fullname == "package.decorator_module.decorator":
            return class_decorator_hook
        return None

def class_decorator_hook(ctx: mypy.plugin.ClassDefContext) -> bool:
    mypy.plugins.common.add_method_to_class(
        ctx.api,
        cls=ctx.cls,
        name="do_assert",
        args=[],  # Instance method with (1 - number of bound params) arguments, i.e. 0 arguments
        return_type=mypy.types.NoneType(),
        self_type=ctx.api.named_type(ctx.cls.fullname),
    )
    del ctx.cls.info.names["do_check"]  # Remove `do_check` from the class
    return True  # Returns whether class is fully defined or needs another round of semantic analysis

test.py

的内容

from package.decorator_module import MyProtocol, decorator

@decorator
class MyClass(MyProtocol):
    def do_check(self) -> bool:
        return False

mc = MyClass()  # mypy: Cannot instantiate abstract class "MyClass" with abstract attribute "do_check" [abstract]
mc.do_check()   # raises `NotImplementedError` at runtime
mc.do_assert()  # OK

package/decorator_module.py

的内容

from __future__ import annotations

import typing_extensions as t

if t.TYPE_CHECKING:
    import collections.abc as cx
    _T = t.TypeVar("_T")

class MyProtocol(t.Protocol):
    def do_check(self) -> bool:
        raise NotImplementedError

# The type annotations here don't mean anything for the mypy plugin,
# which does its own magic when it sees `@package.decorator_module.decorator`.
def decorator(clazz: type[_T]) -> type[_T]:

    do_check: cx.Callable[[_T], bool] = getattr(clazz, "do_check")

    def do_assert(self: _T) -> None:
        assert do_check(self)

    delattr(clazz, "do_check")
    setattr(clazz, "do_assert", do_assert)

    return clazz