Get param/arg map for Python's multi-level Generic's

get_typing_param_args.py

from itertools import zip_longest
from typing import Any, Generic, Literal, Mapping, Protocol, Type, TypeVar, get_args, get_origin, overload


@overload
def get_typing_param_args(
    tp: Type[Any],
    /,
    *,
    resolve: Literal[True],
) -> Mapping[
        TypeVar,
        Type[Any] | None,
]:
    ...


@overload
def get_typing_param_args(
    tp: Type[Any],
    /,
    *,
    resolve: Literal[False],
) -> Mapping[
        TypeVar,
        TypeVar | None,
]:
    ...


def get_typing_param_args(
    type_: Type[Any],
    /,
    *,
    resolve: bool = True,
) -> Mapping[
        TypeVar,
        TypeVar | Type[Any] | None,
]:
    """
    Get the type parameters and their arguments.

    Args:
        tp: The type.
        resolve: Whether to resolve the type parameters into concrete types.

    Returns:
        The type parameters and their arguments.
    """
    param_args: dict[TypeVar, TypeVar | Type[Any] | None] = {}

    types = [type_]
    while types:
        type_ = types.pop(0)
        if (origin := get_origin(type_)) is None or origin in (Generic, Protocol):
            origin = type_
        if (type_params := getattr(origin, "__parameters__", None)) is not None:
            args = get_args(type_)
            for type_param, arg in zip_longest(type_params, args):
                if type_param not in param_args:
                    param_args[type_param] = arg

        if orig_bases := getattr(origin, "__orig_bases__", None):
            types.extend(orig_bases)

    if not resolve:
        return param_args

    for type_param, arg in list(param_args.items()):
        while arg is not None and isinstance(arg, TypeVar):
            if arg is not type_param:
                arg = param_args.get(arg, None)
            elif (default := getattr(type_param, "__default__", None)) is not None \
                    and getattr(default, "has_default", lambda: False)():
                arg = default
            else:
                arg = None

        param_args[type_param] = arg

    return param_args