Skip to content

Commit c71fef0

Browse files
[mypyc] feat: support constant folding in translate_str_format (#19971)
This PR adds support for constant folding inside of `translate_str_format`
1 parent 41efd21 commit c71fef0

File tree

1 file changed

+13
-15
lines changed

1 file changed

+13
-15
lines changed

mypyc/irbuild/specialize.py

Lines changed: 13 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@
7878
uint8_rprimitive,
7979
)
8080
from mypyc.irbuild.builder import IRBuilder
81+
from mypyc.irbuild.constant_fold import constant_fold_expr
8182
from mypyc.irbuild.for_helpers import (
8283
comprehension_helper,
8384
sequence_from_generator_preallocate_helper,
@@ -716,21 +717,18 @@ def translate_dict_setdefault(builder: IRBuilder, expr: CallExpr, callee: RefExp
716717

717718
@specialize_function("format", str_rprimitive)
718719
def translate_str_format(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
719-
if (
720-
isinstance(callee, MemberExpr)
721-
and isinstance(callee.expr, StrExpr)
722-
and expr.arg_kinds.count(ARG_POS) == len(expr.arg_kinds)
723-
):
724-
format_str = callee.expr.value
725-
tokens = tokenizer_format_call(format_str)
726-
if tokens is None:
727-
return None
728-
literals, format_ops = tokens
729-
# Convert variables to strings
730-
substitutions = convert_format_expr_to_str(builder, format_ops, expr.args, expr.line)
731-
if substitutions is None:
732-
return None
733-
return join_formatted_strings(builder, literals, substitutions, expr.line)
720+
if isinstance(callee, MemberExpr):
721+
folded_callee = constant_fold_expr(builder, callee.expr)
722+
if isinstance(folded_callee, str) and expr.arg_kinds.count(ARG_POS) == len(expr.arg_kinds):
723+
tokens = tokenizer_format_call(folded_callee)
724+
if tokens is None:
725+
return None
726+
literals, format_ops = tokens
727+
# Convert variables to strings
728+
substitutions = convert_format_expr_to_str(builder, format_ops, expr.args, expr.line)
729+
if substitutions is None:
730+
return None
731+
return join_formatted_strings(builder, literals, substitutions, expr.line)
734732
return None
735733

736734

0 commit comments

Comments
 (0)