[mypyc]: support :s and :d in generate_format_ops

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
@@ -52,9 +53,9 @@ def generate_format_ops(specifiers: list[ConversionSpecifier]) -> list[FormatOp]
     format_ops = []
     for spec in specifiers:
         # TODO: Match specifiers instead of using whole_seq
-        if spec.whole_seq == "%s" or spec.whole_seq == "{:{}}":
+        if spec.whole_seq == "%s" or spec.whole_seq in ("{:{}}", ":s"):
             format_op = FormatOp.STR
-        elif spec.whole_seq == "%d":
+        elif spec.whole_seq in ("%d", ":d"):
             format_op = FormatOp.INT
         elif spec.whole_seq == "%b":
             format_op = FormatOp.BYTES
@@ -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,50 @@ 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 = "{} {}"
+FMT_D: Final = "{:d} {:s}"
+FMT_S: Final = "{:s} {:s}"
+
+def f() -> str:
+    return FMT.format(400 + 20, "roll" + "up")
+def g() -> str:
+    return FMT_D.format(400 + 20, "roll" + "up")
+def h() -> str:
+    return FMT_S.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
+
+def g():
+    r0, r1, r2, r3 :: str
+L0:
+    r0 = '420'
+    r1 = 'rollup'
+    r2 = ' '
+    r3 = CPyStr_Build(3, r0, r2, r1)
+    return r3
+
+def h():
+    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