Repository: quora/pyanalyze
Pull request #628 (2023-05-23 19:23:12)

Fix use of Literal with new typing_extensions


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 7af71e7..2afdb22 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -1,5 +1,9 @@
 # Changelog
 
+## Unreleased
+
+- Fix use of `Literal` types with `typing_extensions` 4.6.0 (#628)
+
 ## Version 0.10.1 (May 22, 2023)
 
 - Fix errors with protocol matching on `typing_extensions` 4.6.0
diff --git a/pyanalyze/annotations.py b/pyanalyze/annotations.py
index fc162b0..330b203 100644
--- a/pyanalyze/annotations.py
+++ b/pyanalyze/annotations.py
@@ -50,7 +50,7 @@ from typing import (
 import qcore
 
 import typing_inspect
-from typing_extensions import ParamSpec, TypedDict
+from typing_extensions import ParamSpec, TypedDict, get_origin, get_args
 
 from . import type_evaluation
 
@@ -113,15 +113,12 @@ if TYPE_CHECKING:
 
 try:
     from types import GenericAlias
-    from typing import get_args, get_origin  # Python 3.9
 except ImportError:
     GenericAlias = None
-
-    def get_origin(obj: object) -> Any:
-        return None
-
-    def get_args(obj: object) -> Tuple[Any, ...]:
-        return ()
+try:
+    from types import UnionType
+except ImportError:
+    UnionType = None
 
 
 CONTEXT_MANAGER_TYPES = (typing.ContextManager, contextlib.AbstractContextManager)
@@ -418,17 +415,20 @@ def _type_from_runtime(
         return _value_of_origin_args(
             origin, args, val, ctx, allow_unpack=origin is tuple
         )
-    elif typing_inspect.is_literal_type(val):
-        args = typing_inspect.get_args(val)
-        if len(args) == 0:
+    origin = get_origin(val)
+    if is_typing_name(origin, "Literal"):
+        args = get_args(val)
+        if len(args) == 1:
             return KnownValue(args[0])
         else:
             return unite_values(*[KnownValue(arg) for arg in args])
-    elif typing_inspect.is_union_type(val):
-        args = typing_inspect.get_args(val)
+    elif is_typing_name(origin, "Union") or (
+        UnionType is not None and origin is UnionType
+    ):
+        args = get_args(val)
         return unite_values(*[_type_from_runtime(arg, ctx) for arg in args])
-    elif typing_inspect.is_tuple_type(val):
-        args = typing_inspect.get_args(val)
+    elif origin is tuple or is_typing_name(origin, "Tuple"):
+        args = get_args(val)
         if not args:
             if val is tuple or val is Tuple:
                 return TypedValue(tuple)
@@ -490,8 +490,8 @@ def _type_from_runtime(
             is_typeddict=is_typeddict,
             allow_unpack=allow_unpack or origin is tuple or origin is Tuple,
         )
-    elif typing_inspect.is_callable_type(val):
-        args = typing_inspect.get_args(val)
+    elif origin is Callable or is_typing_name(origin, "Callable"):
+        args = get_args(val)
         return _value_of_origin_args(Callable, args, val, ctx)
     elif val is AsynqCallable:
         return CallableValue(Signature.make([ELLIPSIS_PARAM], is_asynq=True))
@@ -511,13 +511,14 @@ def _type_from_runtime(
         if isinstance(val.__supertype__, type):
             # NewType
             return NewTypeValue(val)
-        elif typing_inspect.is_tuple_type(val.__supertype__):
+        super_origin = get_origin(val.__supertype__)
+        if super_origin is tuple or is_typing_name(super_origin, "Tuple"):
             # TODO figure out how to make NewTypes over tuples work
             return AnyValue(AnySource.inference)
         else:
             ctx.show_error(f"Invalid NewType {val}")
             return AnyValue(AnySource.error)
-    elif typing_inspect.is_typevar(val):
+    elif is_typing_name(type(val), "TypeVar"):
         tv = cast(TypeVar, val)
         return make_type_var_value(tv, ctx)
     elif is_instance_of_typing_name(val, "ParamSpec"):
@@ -528,7 +529,7 @@ def _type_from_runtime(
         return ParamSpecKwargsValue(val.__origin__)
     elif is_typing_name(val, "Final") or is_typing_name(val, "ClassVar"):
         return AnyValue(AnySource.incomplete_annotation)
-    elif typing_inspect.is_classvar(val) or typing_inspect.is_final_type(val):
+    elif is_typing_name(origin, "ClassVar") or is_typing_name(origin, "Final"):
         typ = val.__args__[0]
         return _type_from_runtime(typ, ctx)
     elif is_instance_of_typing_name(val, "_ForwardRef") or is_instance_of_typing_name(
@@ -571,12 +572,11 @@ def _type_from_runtime(
         return AnyValue(AnySource.incomplete_annotation)
     elif is_typing_name(val, "TypedDict"):
         return KnownValue(TypedDict)
+    elif isinstance(origin, type):
+        return _maybe_typed_value(origin)
+    elif val is NamedTuple:
+        return TypedValue(tuple)
     else:
-        origin = get_origin(val)
-        if isinstance(origin, type):
-            return _maybe_typed_value(origin)
-        elif val is NamedTuple:
-            return TypedValue(tuple)
         ctx.show_error(f"Invalid type annotation {val}")
         return AnyValue(AnySource.error)
 


Test patch:
diff --git a/pyanalyze/test_annotations.py b/pyanalyze/test_annotations.py
index e08fe36..378f073 100644
--- a/pyanalyze/test_annotations.py
+++ b/pyanalyze/test_annotations.py
@@ -129,6 +129,18 @@ class TestAnnotations(TestNameCheckVisitorBase):
             assert_is_value(x, KnownValue(True))
             assert_is_value(y, MultiValuedValue([KnownValue(True), KnownValue(False)]))
 
+    @assert_passes()
+    def test_literal_in_union(self):
+        from typing import Union
+
+        try:
+            from typing import Literal
+        except ImportError:
+            from typing_extensions import Literal
+
+        def capybara(x: Union[int, Literal["epoch"], None]) -> None:
+            assert_is_value(x, TypedValue(int) | KnownValue("epoch") | KnownValue(None))
+
     @assert_passes()
     def test_contextmanager(self):
         from contextlib import contextmanager