Support for specifying type parameters at function call time (like func[int]()) with runtime introspection

[jonathanslenders]

tl;dr: proof-of-concept decorator at the bottom to make func[int]() work with runtime support.

This has been discussed before:

• Specifying type for Generic Function #943
• The mypy doc mentions: "When calling a generic function, you can’t explicitly pass the values of type parameters as type arguments. The values of type parameters are always inferred by mypy. This is not valid:" https://mypy.readthedocs.io/en/stable/generics.html

We have a justification from the Pydantic world where runtime support is important. Imagine this:

class Model[T](BaseModel):  
    data: T
    
def create_model[T](data: T) -> Model[T]:
    Return Model[T](data=data)

In this case, the [T] is not really needed for the type checkers, because they can infer it from the data argument. But Pydantic needs to know the type parameters at runtime in order to properly deserialize/serialize. So, within the function, we have to be able to resolve the runtime value.

The Mypy docs suggests to use a generic class with a call __call__, however I think that would require an extra pair of parentheses. Anyio has a creative solution by turning the function into a class that inherits from a typed tuple with a __new__, but that only works for functions that return multiple arguments.
https://github.com/agronholm/anyio/blob/master/src/anyio/_core/_streams.py#L16

The anyio approach however doesn't provide runtime access to the value of T due to __orig_class__ not being available in __new__ (for doing get_args(self.__orig_class__)).

What I came up with is a use of __class__getitem__ and creating a new FunctionType where the cells in the __closure__ are replaced with the actual types.

def generic_function[F](func: F) -> F:
    """
    Function decorator that allows calling a generic function with type
    parameters, and expose the actual types within the function.
    """
    class wrapper[T]:
        def __class_getitem__(cls, type_params):
            # Map typevar names to types that we receive in vars.
            typevar_name_to_type = {
                str(name): type for type, name in zip(type_params, func.__type_params__)
            }

            # Helper for creating a function closure.
            def make_cell(value):
                def inner():
                    return value

                return inner.__closure__[0]

            # Create a new closure for the given function by replacing the type
            # variables with the actual types.
            closure = []
            for cell in func.__closure__:
                if isinstance(cell.cell_contents, TypeVar):
                    closure.append(
                        make_cell(typevar_name_to_type[str(cell.cell_contents)])
                    )
                else:
                    closure.append(cell)

            return types.FunctionType(
                func.__code__,
                func.__globals__,
                name=func.__name__,
                argdefs=func.__defaults__,
                closure=tuple(closure),
            )

    return wrapper


@generic_function
def func[T, U](data: T, data2: U) -> None:
    print("Type of T=", T)
    print("Type of U=", U)
    print("value of data=", data, data2)


print(func[int, bool](4, True))

This approach can also be combined with other decorators like contextmanager:
(edit: the following snippet does not work - to make it work with contextmanager we have to dig recursively through the closure and replace the typevars everywhere.)

@generic_function
@contextmanager
def func[T, U](data: T, data2: U) -> None:
    print("Type of T=", T)
    print("Type of U=", U)
    print("value of data=", data, data2)
    yield

with func[int, bool](4, True):
    pass

Is this a reasonable approach? Would it make sense for Python itself to do something similar and substitute the closure?

[jonathanslenders]

Here is an improved implementation that adds support for @contextmanager, with much better runtime performance due to @cache, and better type inference due te the GenericFunction protocol.

from types import FunctionType
from typing import TypeVar, Any, Protocol, Callable
from functools import wraps, cache

__all__ = [
    "GenericFunction",
    "generic_function",
]


class GenericFunction[**P, R](Protocol):
    def __call__(self, *a: P.args, **kw: P.kwargs) -> R: ...
    def __getitem__(self, params: Any) -> Callable[P, R]: ...


def generic_function[**P, R](func: Callable[P, R]) -> GenericFunction[P, R]:
    """
    Function decorator that allows calling a generic function with type
    parameters, and expose the actual types within the function.
    """

    class wrapper:
        @cache
        def __getitem__(self, type_params: TypeVar | tuple[TypeVar, ...]) -> Callable[P, R]:
            if not isinstance(type_params, tuple):
                type_params = (type_params,)
                
            # Map typevar names to types that we receive in vars.
            typevar_name_to_type = {
                name: type_ for type_, name in zip(type_params, func.__type_params__)
            }

            # Helper for creating a function closure.
            def make_cell(value: object) -> Any:
                def inner() -> object:
                    return value

                return inner.__closure__[0]  # type:ignore[index]

            # Create a new closure for the given function by replacing the type
            # variables with the actual types.
            def replace_closure(f: FunctionType) -> tuple[Any, ...]:
                closure = []
                for cell in f.__closure__:  # type:ignore[union-attr]
                    contents = cell.cell_contents
                    if isinstance(contents, TypeVar):
                        closure.append(make_cell(typevar_name_to_type[contents]))
                    elif hasattr(contents, "__closure__"):
                        closure.append(make_cell(replace_func(contents)))
                    else:
                        closure.append(cell)
                return tuple(closure)

            def replace_func(f: FunctionType) -> FunctionType:
                return FunctionType(
                    f.__code__,
                    f.__globals__,
                    name=f.__name__,
                    argdefs=f.__defaults__,
                    closure=replace_closure(f),
                )

            return replace_func(func)  # type:ignore[arg-type]

        # `__call__` staticmethod to make `inspect.signature` work on the
        # `GenericFunction`.
        @staticmethod
        @wraps(func)
        def __call__(*a: P.args, **kw: P.kwargs) -> R:
            return func(*a, **kw)

    wrapper.__doc__ = func.__doc__
    return wrapper()

Example usage:

from contextlib import contextmanager
from typing import Generator
from typing import TYPE_CHECKING
import inspect

@generic_function
def func[T, U](data: T, data2: U) -> tuple[T, U]:
    "Docstring"
    # XXX: Here mypy complains we can't use it T at runtime, but now we can!
    print("Type of T=", T)
    print("Type of U=", U)
    print("value of data=", data, data2)
    return (data, data2)

@generic_function
@contextmanager
def cm_func[T, U](data: T, data2: U) -> Generator[tuple[T, U]]:
    "Generic context manager."
    print("Type of T=", T)
    print("Type of U=", U)
    print("value of data=", data, data2)
    yield (data, data2)

with cm_func[int, bool](4, True) as (a, b):
    if TYPE_CHECKING:
        reveal_type(a)
        reveal_type(b)

c, d = func[int, str](4, "test")
if TYPE_CHECKING:
    reveal_type(c)
    reveal_type(d)
    reveal_type(func[int, str])  # XXX: not yet inferred correctly!

print("Docstrings:")
print("doc=", cm_func.__doc__)
print("doc=", func.__doc__)
print("Signature:")
print(inspect.signature(cm_func))  # Incorrect due to `@contextmanager`, but also without `@generic_function`
print(inspect.signature(func))

Looks like Mypy has some rough edge cases where not everything is well supported, but overall it works great.