How to use overload from the typing package within a higher-order function in Python?

Returning an overloaded function from a higher-order callable does not produce expected results with respect to type hints: namely, the resulting function behaves as if it was unannotated at all.

Here a minimal example:

from typing import Optional, overload


def get_f(b: int):
    @overload
    def f(x: int) -> str: ...

    @overload
    def f(x: int, c: bool) -> int: ...

    def f(x: int, c: Optional[bool] = None) -> int | str:
        return 2 * x * b if c is not None else "s"

    return f


d = get_f(2)
a = d(2)
b = a ** 2

Here mypy does not realize that b = a**2 leads to an TypeError. How to remedy this?

You can create a type denoting an overloaded callable using typing.Protocol (read the PEP to learn more about structural subtyping):

from typing import Optional, Protocol, overload


class FProtocol(Protocol):
    @overload
    def __call__(self, x: int) -> str: ...
    @overload
    def __call__(self, x: int, c: bool) -> int: ...
    

def get_f(b: int) -> FProtocol:
    @overload
    def f(x: int) -> str: ...

    @overload
    def f(x: int, c: bool) -> int: ...

    def f(x: int, c: Optional[bool] = None) -> int | str:
        return 2 * x * b if c is not None else "s"

    return f


d = get_f(2)
a = d(2)
b = a ** 2  # E: Unsupported operand types for ** ("str" and "int")  [operator]

mypy playground pyright playground

Using protocol is the only way to specify an overloaded callable, as far as I know.

When you leave the return type unspecified, it defaults to Any, as PEP484 mandates:

For a checked function, the default annotation for arguments and for the return type is Any