Ensure triple-quoted strings keep working as before

python · Oct 3, 2023 · 472a692 · 472a692
1 parent 2c2c980
commit 472a692
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 14 deletions.
diff --git a/Lib/test/test_tokenize.py b/Lib/test/test_tokenize.py
@@ -577,13 +577,30 @@ def test_string(self):
     NL         '\\n'          (1, 7) (1, 8)
     NAME       'x'           (2, 4) (2, 5)
     OP         ':'           (2, 5) (2, 6)
-    FSTRING_MIDDLE 'a'           (2, 6) (2, 7)
-    NL         '\\n'          (2, 7) (2, 8)
+    FSTRING_MIDDLE 'a\\n'         (2, 6) (3, 0)
     OP         '}'           (3, 0) (3, 1)
     FSTRING_MIDDLE '__'          (3, 1) (3, 3)
     FSTRING_END "'''"         (3, 3) (3, 6)
     """)
         self.check_tokenize("""\
+f'''__{
+    x:a
+    b
+     c
+      d
+}__'''""", """\
+    FSTRING_START "f'''"        (1, 0) (1, 4)
+    FSTRING_MIDDLE '__'          (1, 4) (1, 6)
+    OP         '{'           (1, 6) (1, 7)
+    NL         '\\n'          (1, 7) (1, 8)
+    NAME       'x'           (2, 4) (2, 5)
+    OP         ':'           (2, 5) (2, 6)
+    FSTRING_MIDDLE 'a\\n    b\\n     c\\n      d\\n' (2, 6) (6, 0)
+    OP         '}'           (6, 0) (6, 1)
+    FSTRING_MIDDLE '__'          (6, 1) (6, 3)
+    FSTRING_END "'''"         (6, 3) (6, 6)
+    """)
+        self.check_tokenize("""\
 f'__{
     x:d
 }__'""", """\
@@ -2320,11 +2337,30 @@ def test_string(self):
     LBRACE     '{'           (1, 6) (1, 7)
     NAME       'x'           (2, 4) (2, 5)
     COLON      ':'           (2, 5) (2, 6)
-    FSTRING_MIDDLE 'a'           (2, 6) (2, 7)
+    FSTRING_MIDDLE 'a\\n'         (2, 6) (3, 0)
     RBRACE     '}'           (3, 0) (3, 1)
     FSTRING_MIDDLE '__'          (3, 1) (3, 3)
     FSTRING_END "'''"         (3, 3) (3, 6)
     """)
+
+        self.check_tokenize("""\
+f'''__{
+    x:a
+    b
+     c
+      d
+}__'''""", """\
+    FSTRING_START "f'''"        (1, 0) (1, 4)
+    FSTRING_MIDDLE '__'          (1, 4) (1, 6)
+    LBRACE     '{'           (1, 6) (1, 7)
+    NAME       'x'           (2, 4) (2, 5)
+    COLON      ':'           (2, 5) (2, 6)
+    FSTRING_MIDDLE 'a\\n    b\\n     c\\n      d\\n' (2, 6) (6, 0)
+    RBRACE     '}'           (6, 0) (6, 1)
+    FSTRING_MIDDLE '__'          (6, 1) (6, 3)
+    FSTRING_END "'''"         (6, 3) (6, 6)
+    """)
+
         self.check_tokenize("""\
 f'__{
     x:d

diff --git a/Parser/tokenizer.c b/Parser/tokenizer.c
@@ -2696,21 +2696,22 @@ tok_get_fstring_mode(struct tok_state *tok, tokenizer_mode* current_tok, struct
                 INSIDE_FSTRING_EXPR(current_tok)
         );
 
-        // If we are in a format spec and we found a newline,
-        // it means that the format spec ends here and we should
-        // return to the regular mode.
-        if (in_format_spec && c == '\n') {
-            tok_backup(tok, c);
-            TOK_GET_MODE(tok)->kind = TOK_REGULAR_MODE;
-            p_start = tok->start;
-            p_end = tok->cur;
-            return MAKE_TOKEN(FSTRING_MIDDLE);
-        }
-        if (c == EOF || (current_tok->f_string_quote_size == 1 && c == '\n')) {
+       if (c == EOF || (current_tok->f_string_quote_size == 1 && c == '\n')) {
             if (tok->decoding_erred) {
                 return MAKE_TOKEN(ERRORTOKEN);
             }
 
+            // If we are in a format spec and we found a newline,
+            // it means that the format spec ends here and we should
+            // return to the regular mode.
+            if (in_format_spec && c == '\n') {
+                tok_backup(tok, c);
+                TOK_GET_MODE(tok)->kind = TOK_REGULAR_MODE;
+                p_start = tok->start;
+                p_end = tok->cur;
+                return MAKE_TOKEN(FSTRING_MIDDLE);
+            }
+
             assert(tok->multi_line_start != NULL);
             // shift the tok_state's location into
             // the start of string, and report the error