gh-153569: store f-string state as source spans

Author: pablogsal

Lib/test/test_fstring.py

@@ -1679,6 +1679,10 @@ def __repr__(self):
 
         self.assertEqual(f'{" # nooo "=}', '" # nooo "=\' # nooo \'')
         self.assertEqual(f'{" \" # nooo \" "=}', '" \\" # nooo \\" "=\' " # nooo " \'')
+        self.assertEqual(f'{"""a" # inside"""=}',
+                         '"""a" # inside"""=\'a" # inside\'')
+        self.assertEqual(f"{'''a' # inside'''=}",
+                         "'''a' # inside'''=\"a' # inside\"")
 
         self.assertEqual(f'{ # some comment goes here
   """hello"""=}',  ' \n  """hello"""=\'hello\'')

Lib/test/test_tstring.py

@@ -287,5 +287,9 @@ def test_triple_quoted(self):
         )
         self.assertEqual(fstring(t), "\n        Hello,\n        Python\n        ")
 
+        t = t'{"""a" # inside"""}'
+        self.assertEqual(t.interpolations[0].expression,
+                         '"""a" # inside"""')
+
 if __name__ == '__main__':
     unittest.main()

Parser/lexer/buffer.c

@@ -14,12 +14,6 @@ _PyLexer_SnapshotBuffer(struct tok_state *tok, const char *base,
         ? -1 : tok->line_start - tok->buf;
     snapshot->multi_line_start = tok->multi_line_start == NULL
         ? -1 : tok->multi_line_start - tok->buf;
-    for (int index = tok->tok_mode_stack_index; index >= 0; --index) {
-        tokenizer_mode *mode = &tok->tok_mode_stack[index];
-        mode->start_offset = mode->start == NULL ? -1 : mode->start - tok->buf;
-        mode->multi_line_start_offset = mode->multi_line_start == NULL
-            ? -1 : mode->multi_line_start - tok->buf;
-    }
 }
 
 void
@@ -35,11 +29,4 @@ _PyLexer_RestoreBuffer(struct tok_state *tok, char *base,
         ? NULL : tok->buf + snapshot->line_start;
     tok->multi_line_start = snapshot->multi_line_start < 0
         ? NULL : tok->buf + snapshot->multi_line_start;
-    for (int index = tok->tok_mode_stack_index; index >= 0; --index) {
-        tokenizer_mode *mode = &tok->tok_mode_stack[index];
-        mode->start = mode->start_offset < 0
-            ? NULL : tok->buf + mode->start_offset;
-        mode->multi_line_start = mode->multi_line_start_offset < 0
-            ? NULL : tok->buf + mode->multi_line_start_offset;
-    }
 }

Parser/lexer/lexer.c

@@ -555,8 +555,8 @@ _PyLexer_get_normal_mode(struct tok_state *tok, tokenizer_mode* current_tok, str
          int cursor_in_format_with_debug =
              cursor == 1 && (current_tok->in_debug || in_format_spec);
          int cursor_valid = cursor == 0 || cursor_in_format_with_debug;
-        if ((cursor_valid) && !_PyLexer_update_ftstring_expr(tok, c)) {
-            return MAKE_TOKEN(ENDMARKER);
+        if (cursor_valid) {
+            _PyLexer_update_ftstring_expr(tok, c);
         }
         if ((cursor_valid) && c != '{' && _PyLexer_set_ftstring_expr(tok, token, c)) {
             return MAKE_TOKEN(ERRORTOKEN);

Parser/lexer/lexer.h

@@ -3,10 +3,24 @@
 
 #include "state.h"
 
-int _PyLexer_update_ftstring_expr(struct tok_state *tok, char cur);
-
 int _PyTokenizer_Get(struct tok_state *, struct token *);
 
+static inline const char *
+_PyLexer_SpanView(const struct tok_state *tok, _PyTok_Span span,
+                  Py_ssize_t *length)
+{
+    assert(length != NULL);
+    assert(_PyTok_SpanIsValid(span));
+    assert(tok->buf != NULL);
+    assert(tok->inp >= tok->buf);
+    if (span.start >= tok->buf_offset &&
+            span.end - tok->buf_offset <= tok->inp - tok->buf) {
+        *length = span.end - span.start;
+        return tok->buf + (span.start - tok->buf_offset);
+    }
+    return _PyTok_SourceSpanView(&tok->source, span, length);
+}
+
 /* The view points into the current input window. The next
    _PyTokenizer_Get() call may discard it. */
 static inline const char *
@@ -18,12 +32,7 @@ _PyToken_TextView(const struct tok_state *tok, const struct token *token,
         *length = 0;
         return "";
     }
-    assert(tok->buf != NULL);
-    assert(tok->inp >= tok->buf);
-    assert(token->span.start >= tok->buf_offset);
-    assert(token->span.end - tok->buf_offset <= tok->inp - tok->buf);
-    *length = token->span.end - token->span.start;
-    return tok->buf + (token->span.start - tok->buf_offset);
+    return _PyLexer_SpanView(tok, token->span, length);
 }
 
 #endif

Parser/lexer/lexer_internal.h

@@ -46,6 +46,7 @@ TOK_NEXT_MODE(struct tok_state *tok)
 
 int _PyLexer_nextc(struct tok_state *);
 void _PyLexer_backup(struct tok_state *, int);
+void _PyLexer_update_ftstring_expr(struct tok_state *, char);
 int _PyLexer_set_ftstring_expr(struct tok_state *, struct token *, char);
 int _PyLexer_check_string_prefixes(struct tok_state *, int, int, int, int, int);
 int _PyLexer_scan_number(struct tok_state *, struct token *, int, int);

Parser/lexer/state.c

@@ -60,24 +60,6 @@ _PyTokenizer_tok_new(void)
     return tok;
 }
 
-static void
-free_fstring_expressions(struct tok_state *tok)
-{
-    int index;
-    tokenizer_mode *mode;
-
-    for (index = tok->tok_mode_stack_index; index >= 0; --index) {
-        mode = &(tok->tok_mode_stack[index]);
-        if (mode->last_expr_buffer != NULL) {
-            PyMem_Free(mode->last_expr_buffer);
-            mode->last_expr_buffer = NULL;
-            mode->last_expr_size = 0;
-            mode->last_expr_end = -1;
-            mode->in_format_spec = 0;
-        }
-    }
-}
-
 /* Free a tok_state structure */
 void
 _PyTokenizer_Free(struct tok_state *tok)
@@ -89,7 +71,6 @@ _PyTokenizer_Free(struct tok_state *tok)
     Py_XDECREF(tok->module);
     _PyTok_ReaderFree(tok);
     _PyTok_SourceClear(&tok->source);
-    free_fstring_expressions(tok);
     PyMem_Free(tok);
 }
 

Parser/lexer/state.h

@@ -50,16 +50,11 @@ typedef struct _tokenizer_mode {
     char quote;
     int quote_size;
     int raw;
-    const char* start;
-    const char* multi_line_start;
+    _PyTok_Off start;
+    _PyTok_Off multi_line_start;
     int first_line;
 
-    Py_ssize_t start_offset;
-    Py_ssize_t multi_line_start_offset;
-
-    Py_ssize_t last_expr_size;
-    Py_ssize_t last_expr_end;
-    char* last_expr_buffer;
+    _PyTok_Span debug_expr;
     int in_debug;
     int in_format_spec;