[mypyc] feat: support constant folding in convert_format_expr_to_str [1/1]

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

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

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"