[mypyc] feat: support constant folding in IRBuilder.extract_int [1/1]

mypyc/irbuild/builder.py

@@ -40,7 +40,6 @@
     TypeAlias,
     TypeInfo,
     TypeParam,
-    UnaryExpr,
     Var,
 )
 from mypy.types import (
@@ -106,6 +105,7 @@
     object_rprimitive,
     str_rprimitive,
 )
+from mypyc.irbuild.constant_fold import constant_fold_expr
 from mypyc.irbuild.context import FuncInfo, ImplicitClass
 from mypyc.irbuild.ll_builder import LowLevelIRBuilder
 from mypyc.irbuild.mapper import Mapper
@@ -965,12 +965,8 @@ def maybe_spill_assignable(self, value: Value) -> Register | AssignmentTarget:
         return reg
 
     def extract_int(self, e: Expression) -> int | None:
-        if isinstance(e, IntExpr):
-            return e.value
-        elif isinstance(e, UnaryExpr) and e.op == "-" and isinstance(e.expr, IntExpr):
-            return -e.expr.value
-        else:
-            return None
+        folded = constant_fold_expr(self, e)
+        return folded if isinstance(folded, int) else None
 
     def get_sequence_type(self, expr: Expression) -> RType:
         return self.get_sequence_type_from_type(self.types[expr])

mypyc/irbuild/constant_fold.py

@@ -10,7 +10,7 @@
 
 from __future__ import annotations
 
-from typing import Final, Union
+from typing import TYPE_CHECKING, Final, Union
 
 from mypy.constant_fold import constant_fold_binary_op, constant_fold_unary_op
 from mypy.nodes import (
@@ -26,9 +26,11 @@
     UnaryExpr,
     Var,
 )
-from mypyc.irbuild.builder import IRBuilder
 from mypyc.irbuild.util import bytes_from_str
 
+if TYPE_CHECKING:
+    from mypyc.irbuild.builder import IRBuilder
+
 # All possible result types of constant folding
 ConstantValue = Union[int, float, complex, str, bytes]
 CONST_TYPES: Final = (int, float, complex, str, bytes)