[mypyc] feat: optimize f-string building from Final values

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
+from mypy.nodes import Context, Expression, StrExpr
 from mypy.options import Options
 from mypyc.ir.ops import Integer, Value
 from mypyc.ir.rtypes import (
@@ -143,7 +143,9 @@ 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):
+            if is_str_rprimitive(node_type) or isinstance(
+                x, StrExpr
+            ):  # NOTE: why does mypyc think our fake StrExprs are not str rprimitives?
                 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)

mypyc/irbuild/specialize.py

@@ -31,6 +31,7 @@
     RefExpr,
     StrExpr,
     TupleExpr,
+    Var,
 )
 from mypy.types import AnyType, TypeOfAny
 from mypyc.ir.ops import (
@@ -710,6 +711,28 @@ def translate_fstring(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Va
                 format_ops.append(FormatOp.STR)
                 exprs.append(item.args[0])
 
+        for i in range(len(exprs) - 1):
+
+            def get_literal_str(expr: Expression) -> str | None:
+                # NOTE: not sure where I should put this helper, this file? somewhere else?
+                if isinstance(expr, StrExpr):
+                    return expr.value
+                elif (
+                    isinstance(expr, RefExpr)
+                    and isinstance(expr.node, Var)
+                    and expr.node.final_value is not None
+                ):
+                    return str(expr.node.final_value)
+                return None
+
+            while (
+                len(exprs) >= i + 2
+                and (first := get_literal_str(exprs[i])) is not None
+                and (second := get_literal_str(exprs[i + 1])) is not None
+            ):
+                exprs = [*exprs[:i], StrExpr(first + second), *exprs[i + 2 :]]
+                format_ops = [*format_ops[:i], FormatOp.STR, *format_ops[i + 2 :]]
+
         substitutions = convert_format_expr_to_str(builder, format_ops, exprs, expr.line)
         if substitutions is None:
             return None

mypyc/test-data/irbuild-str.test

@@ -605,3 +605,23 @@ L3:
     r6 = r7
 L4:
     return r6
+
+[case testFStringFromConstants]
+from typing import Final
+string: Final = "abc"
+integer: Final = 123
+floating: Final = 3.14
+boolean: Final = True
+
+def test(x: str) -> str:
+    return f"{string}{integer}{floating}{boolean}{x}{boolean}{floating}{integer}{string}{x}{string}{integer}{floating}{boolean}{x}"
+
+[out]
+def test(x):
+    x, r0, r1, r2, r3 :: str
+L0:
+    r0 = 'abc1233.14True'
+    r1 = 'True3.14123abc'
+    r2 = 'abc1233.14True'
+    r3 = CPyStr_Build(6, r0, x, r1, x, r2, x)
+    return r3