How to annotate a function that returns a callable that should follow a generic Protocol?

[willofferfit]

Say I have a generic protocol Fn[T_contra] that a few callables must follow, and that a user can run a get_fn function that returns the correct callable to the user. The condition for choosing the correct callable is set in configuration Config(NamedTuple).

What should be get_fn's return type, to ensure any typing issues are properly caught? The code below works, but it's using Any, which makes some checks be skipped.

Python 3.12, mypy 1.10.0:

from enum import IntEnum, auto
from typing import Any, NamedTuple, Protocol, TypeVar

T_contra = TypeVar("T_contra", int, str, contravariant=True)


class Fn(Protocol[T_contra]):
    def __call__(self, x: T_contra) -> bool: ...


def fn_int(x: int) -> bool:
    return True


def fn_str(x: str) -> bool:
    return True


def fn_none(x: None) -> bool:
    return True


class FnOption(IntEnum):
    INT = auto()
    STR = auto()


class Config(NamedTuple):
    fn_option: FnOption


def get_fn(config: Config) -> Fn[Any]:  # How to type this properly and concisely?
    match config.fn_option:
        case FnOption.INT:
            return fn_int
        case FnOption.STR:
            return fn_none  # Should fail (fn_str would be correct)
        case _:
            raise AssertionError("Unsupported.")


x = "hi"
fn = get_fn(config=Config(fn_option=FnOption.STR))
y = fn(x=x)  # Should fail the argument's check

If I replace Fn[Any] with Fn[int] | Fn[str] and return fn_none with return fn_str it almost works, but then I get this on line y = fn(x=x):

main.py:44: error: Argument "x" to "__call__" of "Fn" has incompatible type "str"; expected "int" [arg-type]

Which indicates I should probably use @overload. However, shouldn't the type narrowing done by match config.fn_option already take care of that? And if I were to use overload, how would I be able to annotate the input value?

Finally, I'd like to avoid repeating all possible types from T_contra in get_fn's type annotations. Is there a way using generics? E.g. make the return type Fn[T_co] with T_co = TypeVar("T_co", int, str, covariant=True)

[willofferfit]

P.s. GPT-4 doesn't know either

[carljm]

Generics aren't capable of following all the internal logic of get_fn, they just match types in the signature. So in order to do this with generics, you would need T_contra itself to appear directly as the type (or part of the type) of an argument to get_fn. You won't be able to use generics to map from an FnOption value to the corresponding T_contra type argument.

This more general mapping from argument types to return type is exactly what overloads are for, so you could use overloads here if there were an argument of type FnOption to get_fn (and if all call-sites called get_fn with literal values of FnOption)

Unfortunately you introduce a second level of indirection by nesting the FnOption inside the Config. This means you've made the FnOption value effectively invisible to the static type system. With exactly the given runtime code structure, I think the typing you have returning Fn[Any] is the best you'll be able to achieve.

If you're free to adjust the runtime code, you have more options, but it's hard to recommend a specific one without knowing more about your constraints. If you could make Config generic, then you could potentially use generics for this with a signature like def get_fn[T](config: Config[T]) -> Fn[T]: ....

[willofferfit]

Makes sense, thank you so much for the help. In my case changing the structure of the code wouldn't be worth the couple extra type checks, so I'll keep the Any.