Generic MRO

[wyfo]

Generic information are lost with classes __mro__. But it could be useful to have it, especially for #776 implementation. typing module could integrate a generic_mro function in order to retrieve the class MRO annotated with related generic information; this function could also accept a generic alias in argument.

from typing import Generic, TypeVar, generic_mro

T = TypeVar("T")
class A(Generic[T]): ...

U = TypeVar("U")
class B(A[U]): ...

class C(B[int]):
    pass

assert generic_mro(C) == (C, B[int], A[int], Generic, object)
assert generic_mro(B) == (B, A[U], Generic, object)
assert generic_mro(B[int]) == (B[int], A[int], Generic, object)

[wyfo]

Here is a POC of implementation, which could be directly integrated into typing(_extensions)

import sys
from typing import *
from typing import _collect_type_vars

def _generic_mro(result, tp):
    origin = get_origin(tp)
    if origin is None:
        origin = tp
    result[origin] = tp
    if hasattr(origin, "__orig_bases__"):
        parameters = _collect_type_vars(origin.__orig_bases__)
        substitution = dict(zip(parameters, get_args(tp)))
        for base in origin.__orig_bases__:
            if get_origin(base) in result:
                continue
            base_parameters = getattr(base, "__parameters__", ())
            if base_parameters:
                base = base[tuple(substitution.get(p, p) for p in base_parameters)]
            _generic_mro(result, base)

def generic_mro(tp):
    origin = get_origin(tp)
    if origin is None and not hasattr(tp, "__orig_bases__"):
        if not isinstance(tp, type):
            raise TypeError(f"{tp!r} is not a type or a generic alias")
        return tp.__mro__
    # sentinel value to avoid to subscript Generic and Protocol
    result = {Generic: Generic, Protocol: Protocol}
    _generic_mro(result, tp)
    cls = origin if origin is not None else tp
    return tuple(result.get(sub_cls, sub_cls) for sub_cls in cls.__mro__)

This implementation has been tested with Generic but also with PEP 585 builtin generic containters.