How do you type hint inner variable properly in overload

Question

I have a overload function that mypy doesn't like because the return type hint conflicts with the inferred type.

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

I understand why the inferred type is different. I don't know what to do about it.

Is there some way to hint what fn when using an overload like this?

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

Instead of foo returning Union[T_LT, T_LS] you can make it return 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

Or you could hint test_or_script_functions the same way you did the return:

test_or_script_functions: Union[T_LT, T_LS]

But a don't think there's a way to solve it by only changing fn 's hint