[mypyc] Support type narrowing of native int types using "int" (#14524)

JukkaL · web-flow · commit e778a58066f2 · 2023-01-27T11:17:13.000Z
Now `isinstance(x, int)` can be used to narrow a union type that includes a native int type. In mypyc unions there is no runtime distinction between different integer types -- everything is represented at runtime as boxed `int` values anyway. Also test narrowing a native int using the same native int type. Work on mypyc/mypyc#837.
diff --git a/mypy/checker.py b/mypy/checker.py
@@ -178,6 +178,7 @@
 )
 from mypy.types import (
     ANY_STRATEGY,
+    MYPYC_NATIVE_INT_NAMES,
     OVERLOAD_NAMES,
     AnyType,
     BoolTypeQuery,
@@ -4517,10 +4518,7 @@ def analyze_range_native_int_type(self, expr: Expression) -> Type | None:
             ok = True
             for arg in expr.args:
                 argt = get_proper_type(self.lookup_type(arg))
-                if isinstance(argt, Instance) and argt.type.fullname in (
-                    "mypy_extensions.i64",
-                    "mypy_extensions.i32",
-                ):
+                if isinstance(argt, Instance) and argt.type.fullname in MYPYC_NATIVE_INT_NAMES:
                     if native_int is None:
                         native_int = argt
                     elif argt != native_int:
diff --git a/mypy/meet.py b/mypy/meet.py
@@ -15,6 +15,7 @@
 )
 from mypy.typeops import is_recursive_pair, make_simplified_union, tuple_fallback
 from mypy.types import (
+    MYPYC_NATIVE_INT_NAMES,
     AnyType,
     CallableType,
     DeletedType,
@@ -475,6 +476,9 @@ def _type_object_overlap(left: Type, right: Type) -> bool:
         ):
             return True
 
+        if right.type.fullname == "builtins.int" and left.type.fullname in MYPYC_NATIVE_INT_NAMES:
+            return True
+
         # Two unrelated types cannot be partially overlapping: they're disjoint.
         if left.type.has_base(right.type.fullname):
             left = map_instance_to_supertype(left, right.type)
diff --git a/mypy/semanal_classprop.py b/mypy/semanal_classprop.py
@@ -22,7 +22,7 @@
     Var,
 )
 from mypy.options import Options
-from mypy.types import Instance, ProperType
+from mypy.types import MYPYC_NATIVE_INT_NAMES, Instance, ProperType
 
 # Hard coded type promotions (shared between all Python versions).
 # These add extra ad-hoc edges to the subtyping relation. For example,
@@ -177,7 +177,7 @@ def add_type_promotion(
     # Special case the promotions between 'int' and native integer types.
     # These have promotions going both ways, such as from 'int' to 'i64'
     # and 'i64' to 'int', for convenience.
-    if defn.fullname == "mypy_extensions.i64" or defn.fullname == "mypy_extensions.i32":
+    if defn.fullname in MYPYC_NATIVE_INT_NAMES:
         int_sym = builtin_names["int"]
         assert isinstance(int_sym.node, TypeInfo)
         int_sym.node._promote.append(Instance(defn.info, []))
diff --git a/mypy/subtypes.py b/mypy/subtypes.py
@@ -27,6 +27,7 @@
 from mypy.options import Options
 from mypy.state import state
 from mypy.types import (
+    MYPYC_NATIVE_INT_NAMES,
     TUPLE_LIKE_INSTANCE_NAMES,
     TYPED_NAMEDTUPLE_NAMES,
     AnyType,
@@ -1793,14 +1794,19 @@ def covers_at_runtime(item: Type, supertype: Type) -> bool:
         erase_type(item), supertype, ignore_promotions=True, erase_instances=True
     ):
         return True
-    if isinstance(supertype, Instance) and supertype.type.is_protocol:
-        # TODO: Implement more robust support for runtime isinstance() checks, see issue #3827.
-        if is_proper_subtype(item, supertype, ignore_promotions=True):
-            return True
-    if isinstance(item, TypedDictType) and isinstance(supertype, Instance):
-        # Special case useful for selecting TypedDicts from unions using isinstance(x, dict).
-        if supertype.type.fullname == "builtins.dict":
-            return True
+    if isinstance(supertype, Instance):
+        if supertype.type.is_protocol:
+            # TODO: Implement more robust support for runtime isinstance() checks, see issue #3827.
+            if is_proper_subtype(item, supertype, ignore_promotions=True):
+                return True
+        if isinstance(item, TypedDictType):
+            # Special case useful for selecting TypedDicts from unions using isinstance(x, dict).
+            if supertype.type.fullname == "builtins.dict":
+                return True
+        elif isinstance(item, Instance) and supertype.type.fullname == "builtins.int":
+            # "int" covers all native int types
+            if item.type.fullname in MYPYC_NATIVE_INT_NAMES:
+                return True
     # TODO: Add more special cases.
     return False
 
diff --git a/mypy/types.py b/mypy/types.py
@@ -150,6 +150,9 @@
     "typing_extensions.Never",
 )
 
+# Mypyc fixed-width native int types (compatible with builtins.int)
+MYPYC_NATIVE_INT_NAMES: Final = ("mypy_extensions.i64", "mypy_extensions.i32")
+
 DATACLASS_TRANSFORM_NAMES: Final = (
     "typing.dataclass_transform",
     "typing_extensions.dataclass_transform",
diff --git a/mypyc/test-data/irbuild-i64.test b/mypyc/test-data/irbuild-i64.test
@@ -1834,3 +1834,67 @@ L0:
     r0 = CPyLong_FromFloat(x)
     r1 = unbox(int64, r0)
     return r1
+
+[case testI64IsinstanceNarrowing]
+from typing import Union
+from mypy_extensions import i64
+
+class C:
+    a: i64
+
+def narrow1(x: Union[C, i64]) -> i64:
+    if isinstance(x, i64):
+        return x
+    return x.a
+
+def narrow2(x: Union[C, i64]) -> i64:
+    if isinstance(x, int):
+        return x
+    return x.a
+[out]
+def narrow1(x):
+    x :: union[__main__.C, int64]
+    r0 :: object
+    r1 :: int32
+    r2 :: bit
+    r3 :: bool
+    r4 :: int64
+    r5 :: __main__.C
+    r6 :: int64
+L0:
+    r0 = load_address PyLong_Type
+    r1 = PyObject_IsInstance(x, r0)
+    r2 = r1 >= 0 :: signed
+    r3 = truncate r1: int32 to builtins.bool
+    if r3 goto L1 else goto L2 :: bool
+L1:
+    r4 = unbox(int64, x)
+    return r4
+L2:
+    r5 = borrow cast(__main__.C, x)
+    r6 = r5.a
+    keep_alive x
+    return r6
+def narrow2(x):
+    x :: union[__main__.C, int64]
+    r0 :: object
+    r1 :: int32
+    r2 :: bit
+    r3 :: bool
+    r4 :: int64
+    r5 :: __main__.C
+    r6 :: int64
+L0:
+    r0 = load_address PyLong_Type
+    r1 = PyObject_IsInstance(x, r0)
+    r2 = r1 >= 0 :: signed
+    r3 = truncate r1: int32 to builtins.bool
+    if r3 goto L1 else goto L2 :: bool
+L1:
+    r4 = unbox(int64, x)
+    return r4
+L2:
+    r5 = borrow cast(__main__.C, x)
+    r6 = r5.a
+    keep_alive x
+    return r6
diff --git a/mypyc/test-data/run-i64.test b/mypyc/test-data/run-i64.test
@@ -1,5 +1,5 @@
 [case testI64BasicOps]
-from typing import List, Any, Tuple
+from typing import List, Any, Tuple, Union
 
 MYPY = False
 if MYPY:
@@ -497,6 +497,22 @@ def test_for_loop() -> None:
     assert n == 9
     assert sum([x * x for x in range(i64(4 + int()))]) == 1 + 4 + 9
 
+def narrow1(x: Union[str, i64]) -> i64:
+    if isinstance(x, i64):
+        return x
+    return len(x)
+
+def narrow2(x: Union[str, i64]) -> i64:
+    if isinstance(x, int):
+        return x
+    return len(x)
+
+def test_isinstance() -> None:
+    assert narrow1(123) == 123
+    assert narrow1("foobar") == 6
+    assert narrow2(123) == 123
+    assert narrow2("foobar") == 6
+
 [case testI64ErrorValuesAndUndefined]
 from typing import Any, Tuple
 import sys
diff --git a/test-data/unit/check-native-int.test b/test-data/unit/check-native-int.test
@@ -184,3 +184,47 @@ from mypy_extensions import i64, i32
 reveal_type([a for a in range(i64(5))])  # N: Revealed type is "builtins.list[mypy_extensions.i64]"
 [reveal_type(a) for a in range(0, i32(5))]  # N: Revealed type is "mypy_extensions.i32"
 [builtins fixtures/primitives.pyi]
+
+[case testNativeIntNarrowing]
+from typing import Union
+from mypy_extensions import i64, i32
+
+def narrow_i64(x: Union[str, i64]) -> None:
+    if isinstance(x, i64):
+        reveal_type(x)  # N: Revealed type is "mypy_extensions.i64"
+    else:
+        reveal_type(x)  # N: Revealed type is "builtins.str"
+    reveal_type(x)  # N: Revealed type is "Union[builtins.str, mypy_extensions.i64]"
+
+    if isinstance(x, str):
+        reveal_type(x)  # N: Revealed type is "builtins.str"
+    else:
+        reveal_type(x)  # N: Revealed type is "mypy_extensions.i64"
+    reveal_type(x)  # N: Revealed type is "Union[builtins.str, mypy_extensions.i64]"
+
+    if isinstance(x, int):
+        reveal_type(x)  # N: Revealed type is "mypy_extensions.i64"
+    else:
+        reveal_type(x)  # N: Revealed type is "builtins.str"
+    reveal_type(x)  # N: Revealed type is "Union[builtins.str, mypy_extensions.i64]"
+
+def narrow_i32(x: Union[str, i32]) -> None:
+    if isinstance(x, i32):
+        reveal_type(x)  # N: Revealed type is "mypy_extensions.i32"
+    else:
+        reveal_type(x)  # N: Revealed type is "builtins.str"
+    reveal_type(x)  # N: Revealed type is "Union[builtins.str, mypy_extensions.i32]"
+
+    if isinstance(x, str):
+        reveal_type(x)  # N: Revealed type is "builtins.str"
+    else:
+        reveal_type(x)  # N: Revealed type is "mypy_extensions.i32"
+    reveal_type(x)  # N: Revealed type is "Union[builtins.str, mypy_extensions.i32]"
+
+    if isinstance(x, int):
+        reveal_type(x)  # N: Revealed type is "mypy_extensions.i32"
+    else:
+        reveal_type(x)  # N: Revealed type is "builtins.str"
+    reveal_type(x)  # N: Revealed type is "Union[builtins.str, mypy_extensions.i32]"
+
+[builtins fixtures/primitives.pyi]

Original file line number	Diff line number	Diff line change
`@@ -150,6 +150,9 @@`
`150`	`150`	`"typing_extensions.Never",`
`151`	`151`	`)`
`152`	`152`
	`153`	`+# Mypyc fixed-width native int types (compatible with builtins.int)`
	`154`	`+MYPYC_NATIVE_INT_NAMES: Final = ("mypy_extensions.i64", "mypy_extensions.i32")`
	`155`	`+`
`153`	`156`	`DATACLASS_TRANSFORM_NAMES: Final = (`
`154`	`157`	`"typing.dataclass_transform",`
`155`	`158`	`"typing_extensions.dataclass_transform",`