Implement typing.get_orig_class and get_orig_bases

AlexWaygood · Apr 23, 2023 · Feb 11, 2023 · Feb 11, 2023 · Feb 11, 2023 · Feb 11, 2023
commit 50707baa74fd36b1bf02dc05d2510b9ecc5d8a93
diff --git a/Doc/library/typing.rst b/Doc/library/typing.rst
@@ -2328,6 +2328,42 @@ Introspection helpers
 
    .. versionadded:: 3.8
 
+.. function:: get_orig_bases(tp, /)
+
+    Returns the objects in the bases list in the class's definition before 
+    they were modified by ``__mro_entries__``. This is useful for 
+    introspecting ``Generic``\s.
+
+    Examples::
+
+        T = TypeVar("T")
+
+        class Foo(Generic[T]): ...
+
+        class Bar(Foo[int], float): ...
+
+        get_orig_bases(Foo) == (Generic[T],)
+        get_orig_bases(Bar) == (Foo[int], float)
+        get_orig_bases(int) == None
+
+.. function:: get_orig_class(tp, /)
+
+    Returns the ``GenericAlias`` object that was instanciated to create ``tp``.
+
+    Examples::
+
+        T = TypeVar("T")
+
+        class Foo(Generic[T]): ...
+
+        get_orig_class(Foo[int]()) == Foo[int]
+        get_orig_class(list[int]()) == None
+        get_orig_class(int) == None
+
+    .. warning::
+
+        This function will not work inside of the class initalisation process.
+
 .. function:: is_typeddict(tp)
 
    Check if a type is a :class:`TypedDict`.

diff --git a/Lib/test/test_typing.py b/Lib/test/test_typing.py
@@ -4250,6 +4250,23 @@ class C(Generic[T]): pass
                          (Concatenate[int, P], int))
         self.assertEqual(get_args(list | str), (list, str))
 
+    def test_get_orig_class(self):
+        T = TypeVar('T')
+        class C(Generic[T]): pass
+        self.assertEqual(get_orig_class(C[int]()), C[int])
+        self.assertIsNone(get_orig_class(list[int]()))
+        self.assertIsNone(get_orig_class(int))
+
+    def test_get_orig_bases(self):
+        T = TypeVar('T')
+        class C(Generic[T]): pass
+        class D(C[int]): pass
+        class E(C[str], float): pass
+        self.assertEqual(get_orig_bases(C), (Generic[T],))
+        self.assertEqual(get_orig_bases(D), (C[int],))
+        self.assertEqual(get_orig_bases(int), None)
+        self.assertEqual(get_orig_bases(E), (C[str], float))
+
 
 class CollectionsAbcTests(BaseTestCase):
 

diff --git a/Lib/typing.py b/Lib/typing.py
@@ -2273,6 +2273,52 @@ def get_args(tp):
     return ()
 
 
+@overload
+def get_orig_bases(tp: type[object], /) -> tuple[type[Any], ...] | None: ...
+@overload
+def get_orig_bases(tp: Any, /) -> None: ...
+def get_orig_bases(tp: Any, /) -> tuple[type[Any], ...] | None:
+    """Get the __orig_bases__ (see PEP 560) of a class.
+
+    Examples::
+
+        class Foo(Generic[T]): ...
+
+        class Bar(Foo[int]): ...
+
+        get_orig_bases(Foo) == Generic[T]
+        get_orig_bases(Bar) == Foo[int]
+        get_orig_bases(int) == None
+    """
+    if isisinstance(cls, type):
+        try:
+            return cls.__orig_bases__
+        except AttributeError:
+            pass
+    return None
+
+
+def get_orig_class(tp: Any, /) -> GenericAlias | None:
+    """Get the __orig_class__ for an instance of a Generic subclass.
+
+    Examples::
+
+        class Foo(Generic[T]): ...
+
+        get_orig_class(Foo[int]()) == Foo[int]
+        get_orig_class(list[int]()) == None
+        get_orig_class(int) == None
+
+    Warning
+    -------
+    This will always None in the __init__/__new__ methods of the class.
+    """
+    try:
+        return cls.__orig_class__
+    except AttributeError:
+        return None
+
+
 def is_typeddict(tp):
     """Check if an annotation is a TypedDict class