Backport get_origin() and get_args()

jstasiak · jstasiak · commit e4574d11de96 · 2020-02-06T16:07:51.000+01:00
The implementations come from CPython commit 427c84f13f77 with one small change – the get_origin's docstring mentions Annotated as it's also supported. get_origin() and get_args() introduced in [1] and modified in [2] to support Annotated. [1] python/cpython#13685 [2] python/cpython#18260
diff --git a/typing_extensions/src_py3/typing_extensions.py b/typing_extensions/src_py3/typing_extensions.py
@@ -151,7 +151,7 @@ def _check_methods_in_mro(C, *methods):
 HAVE_ANNOTATED = PEP_560 or SUBS_TREE
 
 if PEP_560:
-    __all__.append("get_type_hints")
+    __all__.extend(["get_args", "get_origin", "get_type_hints"])
 
 if HAVE_ANNOTATED:
     __all__.append("Annotated")
@@ -1992,3 +1992,52 @@ class Annotated(metaclass=AnnotatedMeta):
             OptimizedList = Annotated[List[T], runtime.Optimize()]
             OptimizedList[int] == Annotated[List[int], runtime.Optimize()]
         """
+
+# Python 3.9.0+ has those
+if hasattr(typing, 'get_origin'):
+    get_origin = typing.get_origin
+    get_args = typing.get_args
+elif PEP_560:
+    from typing import _GenericAlias  # noqa
+
+    def get_origin(tp):
+        """Get the unsubscripted version of a type.
+
+        This supports generic types, Callable, Tuple, Union, Literal, Final, ClassVar
+        and Annotated. Return None for unsupported types. Examples::
+
+            get_origin(Literal[42]) is Literal
+            get_origin(int) is None
+            get_origin(ClassVar[int]) is ClassVar
+            get_origin(Generic) is Generic
+            get_origin(Generic[T]) is Generic
+            get_origin(Union[T, int]) is Union
+            get_origin(List[Tuple[T, T]][int]) == list
+        """
+        if isinstance(tp, _AnnotatedAlias):
+            return Annotated
+        if isinstance(tp, _GenericAlias):
+            return tp.__origin__
+        if tp is Generic:
+            return Generic
+        return None
+
+    def get_args(tp):
+        """Get type arguments with all substitutions performed.
+
+        For unions, basic simplifications used by Union constructor are performed.
+        Examples::
+            get_args(Dict[str, int]) == (str, int)
+            get_args(int) == ()
+            get_args(Union[int, Union[T, int], str][int]) == (int, str)
+            get_args(Union[int, Tuple[T, int]][str]) == (int, Tuple[str, int])
+            get_args(Callable[[], T][int]) == ([], int)
+        """
+        if isinstance(tp, _AnnotatedAlias):
+            return (tp.__origin__,) + tp.__metadata__
+        if isinstance(tp, _GenericAlias):
+            res = tp.__args__
+            if get_origin(tp) is collections.abc.Callable and res[0] is not Ellipsis:
+                res = (list(res[:-1]), res[-1])
+            return res
+        return ()
