如何在重载中正确键入提示内部变量

Question

我有一个 mypy 不喜欢的重载 function，因为返回类型提示与推断类型冲突。

Incompatible return value type (got "List[Tuple[Union[TestFunc, ScriptFunc], TestDocInfo]]", 
expected "Union[List[Tuple[TestFunc, TestDocInfo]], List[Tuple[ScriptFunc, TestDocInfo]]]")

我明白为什么推断类型不同。 我不知道该怎么办。

在使用这样的重载时，有什么方法可以提示fn是什么吗？

from __future__ import annotations
from typing import List, Type, Union, Tuple, overload
from types import ModuleType

def parse_test_doc_info(fn: Union[TestFunc, ScriptFunc]) -> TestDocInfo:
    ...

class TestDocInfo:
    ...

class TestFunc:
    ...

class ScriptFunc:
    ...

class TestPrefix:
    pat = "test_"

class ScriptPrefix:
    pat = "script_"

T_LT = List[Tuple[TestFunc, TestDocInfo]]
T_LS = List[Tuple[ScriptFunc, TestDocInfo]]

@overload
def foo(module: ModuleType, prefix: Type[TestPrefix] = ...) -> T_LT:
    ...

@overload
def foo(module: ModuleType, prefix: Type[ScriptPrefix]) -> T_LS:
    ...

def foo(module: ModuleType, 
        prefix: Union[Type[TestPrefix], Type[ScriptPrefix]] = TestPrefix
) -> Union[T_LT, T_LS]:
    test_or_script_functions = []
    for name, obj in module.__dict__.items():
        if name.startswith(prefix.pat):
            if callable(obj):
                fn: Union[TestFunc, ScriptFunc] = obj
                # ^ how do you type hint fn without conflicting with return type?
                info: TestDocInfo = parse_test_doc_info(fn)
                test_or_script_functions.append((fn, info))
    return test_or_script_functions  # <-- mypy errors on this

Answer 1

而不是foo返回Union[T_LT, T_LS]你可以让它返回List[Tuple[Union[TestFunc, ScriptFunc], TestDocInfo]] ：

T_LT = List[Tuple[TestFunc, TestDocInfo]]
T_LS = List[Tuple[ScriptFunc, TestDocInfo]]
T_LTS = List[Tuple[Union[TestFunc, ScriptFunc], TestDocInfo]]

@overload
def foo(module: ModuleType, prefix: Type[TestPrefix] = ...) -> T_LT:
    ...

@overload
def foo(module: ModuleType, prefix: Type[ScriptPrefix]) -> T_LS:
    ...

def foo(module: ModuleType, 
        prefix: Union[Type[TestPrefix], Type[ScriptPrefix]] = TestPrefix
) -> T_LTS:
    test_or_script_functions = []
    for name, obj in module.__dict__.items():
        if name.startswith(prefix.pat):
            if callable(obj):
                fn: Union[TestFunc, ScriptFunc] = obj
                info: TestDocInfo = parse_test_doc_info(fn)
                test_or_script_functions.append((fn, info))
    return test_or_script_functions

或者您可以像返回时一样提示test_or_script_functions ：

test_or_script_functions: Union[T_LT, T_LS]

但是不要认为只有改变fn的提示才能解决它