如何在 Python 中表达通用对象的扩展子类型

发布于11月26日

我想在对象上设置一个属性，但保持对象的其余部分不变，例如

from typing import cast, TypeVar, Generic

T = TypeVar("T")


class HasFoo(Generic[T]):
    foo: str


def set_foo_on_obj(obj: T) -> HasFoo[T]:
    setattr(obj, 'foo', 'some_value')
    return cast(HasFoo[T], obj)


def func(a: int) -> int:
    return a


func_with_foo = set_foo_on_obj(func)

在这里，我试图添加属性foo，并告诉类型判断器"它是与以前相同的对象，但现在有一个foo属性.

但上面的示例简单地擦除了可调用对象的其他属性.

推荐答案

正如@paweł-rubin所提到的，只要类型系统中缺少交集类型，就没有一种优雅/直接的方法可以泛化任何给定的类型.

您可以针对特定用例编写具有不同复杂程度的变通方法，但是使用typing.Protocol已经提供的 struct 子类型.如果您想用添加的协议来"增强"简单的Callable，您可以这样做:

from collections.abc import Callable
from typing import Generic, ParamSpec, Protocol, TypeVar, cast


T = TypeVar("T")
P = ParamSpec("P")


class HasFoo(Protocol):
    foo: str


class CallableWithFoo(HasFoo, Generic[P, T]):
    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T: ...


def set_foo_on_func(function: Callable[P, T]) -> CallableWithFoo[P, T]:
    function.foo = "some_value"
    return cast(CallableWithFoo[P, T], function)


def func(a: int) -> int:
    return a


func_with_foo = set_foo_on_func(func)

reveal_type(func_with_foo)      # CallableWithFoo[[a: builtins.int], builtins.int]
reveal_type(func_with_foo(1))   # builtins.int
reveal_type(func_with_foo.foo)  # builtins.str

在我们的泛型类中使用typing.ParamSpec允许在修饰后保留可调用签名.

显然，其他类型(不是Callable个子类型)将需要其他协议继承.但如果没有适当的交叉口类型，这可能就是最好的结果了.

也没有办法避免typing.castIMO，因为由于显而易见的原因，静态类型判断器会忽略动态属性分配.

EDIT:将setattr(function, "foo", "some_value")更改为常规属性赋值.感谢@SuTerliakov指出.