[mypyc] Faster fstring

python · JukkaL · Jul 8, 2021 · Jul 6, 2021 · Jul 7, 2021 · Jul 7, 2021
commit 4ac5f4e90d24a34c59f8aa48f8646c8384cb651a
diff --git a/mypyc/irbuild/specialize.py b/mypyc/irbuild/specialize.py
@@ -451,3 +451,39 @@ def split_braces(format_str: str) -> List[str]:
             prev = ''
     ret_list.append(tmp_str)
     return ret_list
+
+@specialize_function('join', str_rprimitive)
+def translate_fstring(
+        builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Optional[Value]:
+    # Special case for f-string, which is translated into str.join() in mypy AST.
+    # This specializer optimizes simplest f-strings which don't contain any
+    # format operation.
+    if expr.arg_kinds == [ARG_POS] and isinstance(expr.args[0], ListExpr):
+        for item in expr.args[0].items:
+            if isinstance(item, StrExpr):
+                continue
+            elif isinstance(item, CallExpr):
+                if item.callee.name != 'format':
+                    return None
+                if item.callee.expr.value != '{:{}}':
+                    return None
+                if not isinstance(item.args[1], StrExpr) or item.args[1].value != '':
+                    return None
+            else:
+                return None
+
+        result_list: List[Value] = [Integer(0, c_pyssize_t_rprimitive)]
+        for item in expr.args[0].items:
+            if isinstance(item, StrExpr):
+                result_list.append(builder.accept(item))
+            elif isinstance(item, CallExpr):
+                result_list.append(builder.call_c(str_op,
+                                                  [builder.accept(item.args[0])],
+                                                  expr.line))
+
+        if len(result_list) == 1:
+            return builder.load_str("")
+
+        result_list[0] = Integer(len(result_list) - 1, c_pyssize_t_rprimitive)
+        return builder.call_c(str_build_op, result_list, expr.line)
+    return None
diff --git a/mypyc/test-data/irbuild-str.test b/mypyc/test-data/irbuild-str.test
@@ -201,3 +201,58 @@ L0:
     r17 = CPyStr_Build(6, r14, r9, r15, r11, r16, r13)
     s3 = r17
     return 1
+
+[case testFStrings]
+def f(var: str, num: int) -> None:
+    s1 = f"Hi! I'm {var}. I am {num} years old."
+    s2 = f'Hello {var:>{num}}'
+    s3 = f''
+    s4 = f'abc'
+[out]
+def f(var, num):
+    var :: str
+    num :: int
+    r0, r1, r2 :: str
+    r3 :: object
+    r4, r5, r6, s1, r7, r8, r9, r10 :: str
+    r11 :: object
+    r12, r13, r14 :: str
+    r15 :: object
+    r16 :: str
+    r17 :: list
+    r18, r19, r20 :: ptr
+    r21, s2, r22, s3, r23, s4 :: str
+L0:
+    r0 = "Hi! I'm "
+    r1 = PyObject_Str(var)
+    r2 = '. I am '
+    r3 = box(int, num)
+    r4 = PyObject_Str(r3)
+    r5 = ' years old.'
+    r6 = CPyStr_Build(5, r0, r1, r2, r4, r5)
+    s1 = r6
+    r7 = ''
+    r8 = 'Hello '
+    r9 = '{:{}}'
+    r10 = '>'
+    r11 = box(int, num)
+    r12 = PyObject_Str(r11)
+    r13 = CPyStr_Build(2, r10, r12)
+    r14 = 'format'
+    r15 = CPyObject_CallMethodObjArgs(r9, r14, var, r13, 0)
+    r16 = cast(str, r15)
+    r17 = PyList_New(2)
+    r18 = get_element_ptr r17 ob_item :: PyListObject
+    r19 = load_mem r18 :: ptr*
+    set_mem r19, r8 :: builtins.object*
+    r20 = r19 + 8
+    set_mem r20, r16 :: builtins.object*
+    keep_alive r17
+    r21 = PyUnicode_Join(r7, r17)
+    s2 = r21
+    r22 = ''
+    s3 = r22
+    r23 = 'abc'
+    s4 = r23
+    return 1
+