diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py index 125382145991..63930123135f 100644 --- a/mypyc/irbuild/builder.py +++ b/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]) diff --git a/mypyc/irbuild/constant_fold.py b/mypyc/irbuild/constant_fold.py index 12a4b15dd40c..b1133f95b18e 100644 --- a/mypyc/irbuild/constant_fold.py +++ b/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)