Fix :s format handling and tests

Author: BobTheBuidler

mypyc/irbuild/format_str_tokenizer.py

@@ -38,9 +38,11 @@ class FormatOp(Enum):
     For example, to mark a conversion from any object to string,
     ConversionSpecifier may have several representations, like '%s', '{}'
     or '{:{}}'. However, there would only exist one corresponding FormatOp.
+    The ':s' format code is stricter and maps to STR_ONLY.
     """
 
     STR = "s"
+    STR_ONLY = "s_only"
     INT = "d"
     BYTES = "b"
 
@@ -53,8 +55,10 @@ 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 in ("{:{}}", ":s"):
+        if spec.whole_seq == "%s" or spec.whole_seq == "{:{}}":
             format_op = FormatOp.STR
+        elif spec.whole_seq == ":s":
+            format_op = FormatOp.STR_ONLY
         elif spec.whole_seq in ("%d", ":d"):
             format_op = FormatOp.INT
         elif spec.whole_seq == "%b":
@@ -152,6 +156,13 @@ def convert_format_expr_to_str(
                 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.STR_ONLY:
+            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)
+            else:
+                return None
         elif format_op == FormatOp.INT:
             if isinstance(folded := constant_fold_expr(builder, x), int):
                 var_str = builder.load_literal_value(str(folded))

mypyc/test-data/irbuild-constant-fold.test

@@ -200,7 +200,7 @@ def f() -> str:
 def g() -> str:
     return FMT_D.format(400 + 20, "roll" + "up")
 def h() -> str:
-    return FMT_S.format(400 + 20, "roll" + "up")
+    return FMT_S.format("4" + "20", "roll" + "up")
 
 [out]
 def f():
@@ -222,7 +222,7 @@ L0:
 def h():
     r0, r1, r2, r3 :: str
 L0:
-    r0 = CPyTagged_Str(840)
+    r0 = '420'
     r1 = 'rollup'
     r2 = ' '
     r3 = CPyStr_Build(3, r0, r2, r1)