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

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()

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))