diff --git a/mypyc/irbuild/format_str_tokenizer.py b/mypyc/irbuild/format_str_tokenizer.py index 5a35900006d2..2596e8e050a6 100644 --- a/mypyc/irbuild/format_str_tokenizer.py +++ b/mypyc/irbuild/format_str_tokenizer.py @@ -12,7 +12,7 @@ ) from mypy.errors import Errors from mypy.messages import MessageBuilder -from mypy.nodes import Context, Expression, StrExpr +from mypy.nodes import Context, Expression from mypy.options import Options from mypyc.ir.ops import Integer, Value from mypyc.ir.rtypes import ( @@ -23,6 +23,7 @@ is_str_rprimitive, ) from mypyc.irbuild.builder import IRBuilder +from mypyc.irbuild.constant_fold import constant_fold_expr from mypyc.primitives.bytes_ops import bytes_build_op from mypyc.primitives.int_ops import int_to_str_op from mypyc.primitives.str_ops import str_build_op, str_op @@ -143,16 +144,18 @@ def convert_format_expr_to_str( for x, format_op in zip(exprs, format_ops): node_type = builder.node_type(x) if format_op == FormatOp.STR: - if is_str_rprimitive(node_type) or isinstance( - x, StrExpr - ): # NOTE: why does mypyc think our fake StrExprs are not str rprimitives? + if isinstance(folded := constant_fold_expr(builder, x), str): + var_str = builder.load_literal_value(folded) + elif is_str_rprimitive(node_type): var_str = builder.accept(x) elif is_int_rprimitive(node_type) or is_short_int_rprimitive(node_type): var_str = builder.primitive_op(int_to_str_op, [builder.accept(x)], line) else: var_str = builder.primitive_op(str_op, [builder.accept(x)], line) elif format_op == FormatOp.INT: - if is_int_rprimitive(node_type) or is_short_int_rprimitive(node_type): + if isinstance(folded := constant_fold_expr(builder, x), int): + var_str = builder.load_literal_value(str(folded)) + elif is_int_rprimitive(node_type) or is_short_int_rprimitive(node_type): var_str = builder.primitive_op(int_to_str_op, [builder.accept(x)], line) else: return None diff --git a/mypyc/test-data/irbuild-constant-fold.test b/mypyc/test-data/irbuild-constant-fold.test index cd953c84c541..009a4e68f678 100644 --- a/mypyc/test-data/irbuild-constant-fold.test +++ b/mypyc/test-data/irbuild-constant-fold.test @@ -186,6 +186,25 @@ L0: big5 = r4 return 1 +[case testConstantFoldFormatArgs] +# This only tests that the callee and args are constant folded, +# it is not intended to test the result. +from typing import Any, Final + +FMT: Final = "{} {}" + +def f() -> str: + return FMT.format(400 + 20, "roll" + "up") +[out] +def f(): + r0, r1, r2, r3 :: str +L0: + r0 = CPyTagged_Str(840) + r1 = 'rollup' + r2 = ' ' + r3 = CPyStr_Build(3, r0, r2, r1) + return r3 + [case testIntConstantFoldingFinal] from typing import Final X: Final = 5 diff --git a/mypyc/test-data/run-strings.test b/mypyc/test-data/run-strings.test index 0ae67ed7f1c3..ea1015596ad5 100644 --- a/mypyc/test-data/run-strings.test +++ b/mypyc/test-data/run-strings.test @@ -1137,3 +1137,11 @@ def test_optional_ne() -> None: assert ne_s_opt_s_opt('y', 'x') assert ne_s_opt_s_opt(None, 'x') assert ne_s_opt_s_opt('x', None) + +[case testConstantFoldFormatArgs] +from typing import Final + +FMT: Final = "{} {}" + +def test_format() -> str: + assert FMT.format(400 + 20, "roll" + "up") == "420 rollup"