From 5c7d1028fc8f8f94d34955627f9cfd284110a84a Mon Sep 17 00:00:00 2001
From: Jelle Zijlstra <jelle.zijlstra@gmail.com>
Date: Sun, 30 Jan 2022 18:37:36 -0800
Subject: [PATCH] Treat Thrift enums as compatible with int (#455)

---
 docs/changelog.md             | 1 +
 pyanalyze/test_thrift_enum.py | 5 +++++
 pyanalyze/type_object.py      | 2 ++
 3 files changed, 8 insertions(+)

diff --git a/docs/changelog.md b/docs/changelog.md
index 8c51c023..a61a44b2 100644
--- 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
index e31647ec..c0e393f5 100644
--- 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
index 904d1327..2c56765d 100644
--- 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: