From db5b7efc9da84d292884cbec5370331c0bb760db Mon Sep 17 00:00:00 2001
From: Jelle Zijlstra <jelle.zijlstra@gmail.com>
Date: Sat, 2 Mar 2024 07:35:08 -0800
Subject: [PATCH] Apply implicit int/float promotion to subclasses (#738)

Fixes #736
---
 docs/changelog.md             |  1 +
 pyanalyze/test_type_object.py | 13 +++++++------
 pyanalyze/type_object.py      |  4 ++--
 3 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/docs/changelog.md b/docs/changelog.md
index 53499018..480eaf5d 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Treat subclasses of `int` as subclasses of `float` and `complex` too (#738)
 - Fix crash on encountering certain decorators in stubs (#734)
 - Fix inference of signature for certain secondary methods (#732)
 
diff --git a/pyanalyze/test_type_object.py b/pyanalyze/test_type_object.py
index 0489ba85..6cc4d6b3 100644
--- a/pyanalyze/test_type_object.py
+++ b/pyanalyze/test_type_object.py
@@ -23,6 +23,9 @@ def test_float(self):
         def take_float(x: float) -> None:
             pass
 
+        class IntSubclass(int):
+            pass
+
         def capybara(nt: NT, i: int, f: float) -> None:
             take_float(nt)
             take_float(i)
@@ -31,11 +34,9 @@ def capybara(nt: NT, i: int, f: float) -> None:
             take_float(3)
             take_float(1 + 1j)  # E: incompatible_argument
             take_float("string")  # E: incompatible_argument
-            # Strictly speaking this should be allowed, but it doesn't
-            # seem useful and I don't know of a concrete use case that would
-            # require this to work. This can be revisited if we do find a use
-            # case.
-            take_float(True)  # E: incompatible_argument
+            # bool is a subclass of int, which is treated as a subclass of float
+            take_float(True)
+            take_float(IntSubclass(3))
 
     @assert_passes()
     def test_complex(self):
@@ -57,7 +58,7 @@ def capybara(nti: NTI, ntf: NTF, i: int, f: float, c: complex) -> None:
             take_complex(3)
             take_complex(1 + 1j)
             take_complex("string")  # E: incompatible_argument
-            take_complex(True)  # E: incompatible_argument
+            take_complex(True)  # bool is an int, which is a float, which is a complex
 
 
 class TestSyntheticType(TestNameCheckVisitorBase):
diff --git a/pyanalyze/type_object.py b/pyanalyze/type_object.py
index edff8faa..4c06247e 100644
--- a/pyanalyze/type_object.py
+++ b/pyanalyze/type_object.py
@@ -68,10 +68,10 @@ def __post_init__(self) -> None:
         self.base_classes |= set(get_mro(self.typ))
         # As a special case, the Python type system treats int as
         # a subtype of float, and both int and float as subtypes of complex.
-        if self.typ is int:
+        if self.typ is int or safe_in(int, self.base_classes):
             self.artificial_bases.add(float)
             self.artificial_bases.add(complex)
-        if self.typ is float:
+        if self.typ is float or safe_in(float, self.base_classes):
             self.artificial_bases.add(complex)
         if self.is_thrift_enum:
             self.artificial_bases.add(int)