Treat Thrift enums as compatible with int (#455)

quora · Jan 31, 2022 · 5c7d102 · 5c7d102
1 parent c0ebfd4
commit 5c7d102
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 0 deletions.
diff --git a/docs/changelog.md b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Treat Thrift enums as compatible with `int` (#455)
 - Fix treatment of `TypeVar` with bounds or constraints
   as callables (#454)
 - Improve `TypeVar` solution algorithm (#453)

diff --git a/pyanalyze/test_thrift_enum.py b/pyanalyze/test_thrift_enum.py
@@ -18,6 +18,9 @@ class ThriftEnum(object):
         def want_enum(e: ThriftEnum):
             pass
 
+        def want_int(i: int):
+            pass
+
         def capybara(e: ThriftEnum):
             want_enum(e)
             want_enum(ThriftEnum.X)
@@ -26,6 +29,8 @@ def capybara(e: ThriftEnum):
             want_enum(1)
             want_enum(42)  # E: incompatible_argument
             want_enum(str(e))  # E: incompatible_argument
+            want_int(e)
+            want_int(e.X)
 
     @assert_passes()
     def test_typevar(self):

diff --git a/pyanalyze/type_object.py b/pyanalyze/type_object.py
@@ -67,6 +67,8 @@ def __post_init__(self) -> None:
             self.base_classes.add(complex)
         if self.typ is float:
             self.base_classes.add(complex)
+        if self.is_thrift_enum:
+            self.base_classes.add(int)
 
     def is_assignable_to_type(self, typ: type) -> bool:
         for base in self.base_classes: