Repository: quora/pyanalyze
Pull request #394 (2022-01-11 00:34:47)

fix Thrift enum + TypeVar


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 39d5168..c0bf0bb 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Fix compatibility between Thrift enums and TypeVars (#394)
 - Fix accessing attributes on Unions nested within
   Annotated (#393)
 - Fix interaction of `register_error_code()` with new
diff --git a/pyanalyze/value.py b/pyanalyze/value.py
index 4c548be..afbbccf 100644
--- a/pyanalyze/value.py
+++ b/pyanalyze/value.py
@@ -619,6 +619,8 @@ class TypedValue(Value):
         if isinstance(other, AnyValue) and not ctx.should_exclude_any():
             ctx.record_any_used()
             return {}
+        elif isinstance(other, TypeVarValue):
+            return super().can_assign(other, ctx)
         elif isinstance(other, KnownValue):
             if not isinstance(other.val, int):
                 return CanAssignError(f"{other} is not an int")
@@ -1483,7 +1485,7 @@ class TypeVarValue(Value):
         return typevars.get(self.typevar, self)
 
     def can_assign(self, other: Value, ctx: CanAssignContext) -> CanAssign:
-        if self == other:
+        if self == other or (isinstance(other, AnnotatedValue) and self == other.value):
             return {}
         if self.bound is not None:
             can_assign = self.bound.can_assign(other, ctx)


Test patch:
diff --git a/pyanalyze/test_thrift_enum.py b/pyanalyze/test_thrift_enum.py
new file mode 100644
index 0000000..e31647e
--- /dev/null
+++ b/pyanalyze/test_thrift_enum.py
@@ -0,0 +1,60 @@
+# static analysis: ignore
+
+from .test_name_check_visitor import TestNameCheckVisitorBase
+from .test_node_visitor import assert_passes
+from .value import KnownValue, TypedValue, assert_is_value
+
+
+class TestThriftEnum(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test_basic(self):
+        class ThriftEnum(object):
+            X = 0
+            Y = 1
+
+            _VALUES_TO_NAMES = {0: "X", 1: "Y"}
+            _NAMES_TO_VALUES = {"X": 0, "Y": 1}
+
+        def want_enum(e: ThriftEnum):
+            pass
+
+        def capybara(e: ThriftEnum):
+            want_enum(e)
+            want_enum(ThriftEnum.X)
+            want_enum(ThriftEnum.Y)
+            want_enum(0)
+            want_enum(1)
+            want_enum(42)  # E: incompatible_argument
+            want_enum(str(e))  # E: incompatible_argument
+
+    @assert_passes()
+    def test_typevar(self):
+        from typing import TypeVar
+        from typing_extensions import Annotated
+
+        class ThriftEnum(object):
+            X = 0
+            Y = 1
+
+            _VALUES_TO_NAMES = {0: "X", 1: "Y"}
+            _NAMES_TO_VALUES = {"X": 0, "Y": 1}
+
+        TET = TypeVar("TET", bound=ThriftEnum)
+
+        def want_enum(te: ThriftEnum) -> None:
+            pass
+
+        def get_it(te: TET) -> TET:
+            want_enum(te)
+            return te
+
+        def get_it_annotated(te: Annotated[TET, 3]) -> TET:
+            want_enum(te)
+            return te
+
+        def capybara(e: ThriftEnum):
+            assert_is_value(get_it(e), TypedValue(ThriftEnum))
+            assert_is_value(get_it(ThriftEnum.X), KnownValue(ThriftEnum.X))
+
+            assert_is_value(get_it_annotated(e), TypedValue(ThriftEnum))
+            assert_is_value(get_it_annotated(ThriftEnum.X), KnownValue(ThriftEnum.X))