feat: implement constant folding for str.format and f-strings [2/3]

mypy/constant_fold.py

@@ -8,13 +8,17 @@
 from typing import Final, Union
 
 from mypy.nodes import (
+    CallExpr,
     ComplexExpr,
     Expression,
     FloatExpr,
     IntExpr,
+    ListExpr,
+    MemberExpr,
     NameExpr,
     OpExpr,
     StrExpr,
+    TupleExpr,
     UnaryExpr,
     Var,
 )
@@ -73,6 +77,34 @@ def constant_fold_expr(expr: Expression, cur_mod_id: str) -> ConstantValue | Non
         value = constant_fold_expr(expr.expr, cur_mod_id)
         if value is not None:
             return constant_fold_unary_op(expr.op, value)
+    # --- f-string requires partial support for both str.join and str.format ---
+    elif (
+        isinstance(expr, CallExpr)
+        and isinstance(callee := expr.callee, MemberExpr)
+        and isinstance(folded_callee := constant_fold_expr(callee.expr, cur_mod_id), str)
+    ):
+        # --- partial str.join constant folding ---
+        if (
+            callee.name == "join"
+            and len(args := expr.args) == 1
+            and isinstance(arg := args[0], (ListExpr, TupleExpr))
+        ):
+            folded_items: list[str] = []
+            for item in arg.items:
+                val = constant_fold_expr(item, cur_mod_id)
+                if not isinstance(val, str):
+                    return None
+                folded_items.append(val)
+            return folded_callee.join(folded_items)
+        # --- str.format constant folding
+        elif callee.name == "format":
+            folded_args: list[ConstantValue] = []
+            for arg in expr.args:
+                arg_val = constant_fold_expr(arg, cur_mod_id)
+                if arg_val is None:
+                    return None
+                folded_args.append(arg_val)
+            return folded_callee.format(*folded_args)
     return None
 
 

mypyc/irbuild/constant_fold.py

@@ -15,6 +15,7 @@
 from mypy.constant_fold import constant_fold_binary_op, constant_fold_unary_op
 from mypy.nodes import (
     BytesExpr,
+    CallExpr,
     ComplexExpr,
     Expression,
     FloatExpr,
@@ -74,6 +75,20 @@ def constant_fold_expr(builder: IRBuilder, expr: Expression) -> ConstantValue |
         value = constant_fold_expr(builder, expr.expr)
         if value is not None and not isinstance(value, bytes):
             return constant_fold_unary_op(expr.op, value)
+    # --- str.format constant folding
+    elif (
+        isinstance(expr, CallExpr)
+        and isinstance(callee := expr.callee, MemberExpr)
+        and callee.name == "format"
+        and isinstance(folded_callee := constant_fold_expr(builder, callee), str)
+    ):
+        folded_args: list[ConstantValue] = []
+        for arg in expr.args:
+            arg_val = constant_fold_expr(builder, arg)
+            if arg_val is None:
+                return None
+            folded_args.append(arg_val)
+        return folded_callee.format(*folded_args)
     return None