mypy 从被调用函数推断类型签名

有没有办法从被调用函数推断调用者的类型签名？

在下面的示例中，（重载） Test1.f 函数按预期工作。 Test2.get_arg().f 也有效。

问题是：以类型安全的方式实现 Test2.f 的正确（最短）方法是什么，或者最接近的解决方法是什么？ Test2.f 只是 Test2.get_arg().f 的快捷方式。如果可能的话，我想避免重复 Test2.f 的类型签名。

这是示例：

from typing import *

class Test1:
    @overload
    def f(self, key: Literal['A']) -> int: ...
    @overload
    def f(self, key: Literal['B']) -> str: ...
    def f(self, key: Any) -> Any:
        if key == 'A': return 1
        elif key == 'B': return 'x'
        assert_never(key)

class Test2:
    def __init__(self, arg: Test1):
        self._arg = arg

    def get_arg(self) -> Test1:
        return self._arg

    # How to tell to 'mypy' that 'f' is the same as 'get_arg().f',
    # without repeating the (overload) definition?
    # This works in runtime, but not for 'mypy'
    def f(self, key):
        return self.get_arg().f(key)

test2 = Test2(Test1())

# this works as expected
print(test2.get_arg().f('A') + 1)
#print(test2.get_arg().f('B') + 1) # <- mypy error (str + int)
#print(test2.get_arg().f('C')) # <- mypy error (no overload)

# we would like to have the same type safety with shortcut 'f'
print(test2.f('A') + 1)
#print(test2.f('B') + 1) # <- we want type error here
#print(test2.f('C')) # <- and here

我尝试过放置各种类型签名，例如：

def f(self, key: Any) -> Any:
def f(self, key: Union[Literal['A'], Literal['B']]) -> Any:

它修复了一些问题，但是 mypy 仍然不知道返回类型是什么。我不确定这种做法是否正确。

很遗憾，在不重复 Test1.f 的类型签名的情况下，你无法以 mypy 能够理解的方式直接表达 Test2.f 与 Test2.get_arg().f 相同的事实。

这里的问题是 mypy 没有进行足够深入的代码流分析来理解 Test2.f 仅仅是委托给 self.get_arg().f 。

以下是几种解决方法，每种方法都有其优缺点：

1. 使用 cast

def f(self, key):
    return cast(Any, self.get_arg().f(key))  # type: ignore

• 优点： 简单直接。
• 缺点： 放弃了类型安全， mypy 不会捕获 Test2.f 调用中的类型错误。

2. 使用泛型

T = TypeVar('T')

class Test2:
    ...

    def f(self, key: str) -> T:
        return self.get_arg().f(key)  # type: ignore

• 优点： 保留了一些类型安全， mypy 会根据调用 Test2.f 的位置推断返回类型。
• 缺点： 仍然需要使用 type: ignore ，因为 mypy 不能保证 self.get_arg().f(key) 的返回值总是与 T 匹配。

3. 重复类型签名

    @overload
    def f(self, key: Literal['A']) -> int: ...
    @overload
    def f(self, key: Literal['B']) -> str: ...
    def f(self, key: Any) -> Any:
        return self.get_arg().f(key)

• 优点： 提供最佳的类型安全， mypy 会完全理解 Test2.f 的类型行为。
• 缺点： 需要重复代码，违背了你想要避免的目标。

4. 使用组合而不是继承

class Test2:
    def __init__(self, arg: Test1):
        self.arg = arg

    def f(self, key: str) -> Any:
        return self.arg.f(key)

• 优点： 避免了重复类型签名，并保持了类型安全。
• 缺点： 需要修改类的结构，可能不适用于所有情况。

哪种方法最适合你取决于你项目的具体情况和你对类型安全的容忍度。如果你绝对想要避免代码重复并且可以接受一些类型安全的损失，那么 cast 或泛型可能是合适的选择。如果你需要最高级别的类型安全，那么重复类型签名或使用组合是更好的选择。