gh-153569: return token spans and unify decoded source storage (#156482)

Represent token boundaries as logical source offsets with explicit start
and end locations, and adapt pegen and _tokenize to consume them.

Use SourceText to own decoded input across all reader modes. Streaming
readers discard consumed windows while preserving logical offsets and
reusing the allocation. Prepared and interactive input retain their source.

Force source relocation on growth in debug builds and add coverage for
formatted-string growth, interactive input, storage reuse, and offset
limits. Scanner positions remain pointers, rebased when storage moves.

Author: pablogsal

Lib/test/test_capi/test_tokenizer.py

@@ -9,6 +9,9 @@ class TokenizerTests(unittest.TestCase):
     def test_source(self):
         _testinternalcapi.test_tokenizer_source()
 
+    def test_source_discard(self):
+        _testinternalcapi.test_tokenizer_source_discard()
+
     def test_cursor(self):
         _testinternalcapi.test_tokenizer_cursor()
 

Lib/test/test_repl.py

@@ -198,6 +198,22 @@ def test_lexer_buffer_realloc_with_null_start(self):
         self.assertEqual(p.returncode, 0)
         self.assertIn(long_value, output)
 
+    @cpython_only
+    def test_multiline_fstring_source_reallocation(self):
+        long_line = " " * 9000 + "+ 2"
+        user_input = (
+            'value = f"""{(\n'
+            '1\n'
+            f'{long_line}\n'
+            ')}"""\n'
+            'print(value)\n'
+        )
+        p = spawn_repl()
+        p.stdin.write(user_input)
+        output = kill_python(p)
+        self.assertEqual(p.returncode, 0)
+        self.assertIn(">>> 3\n>>> ", output)
+
     def test_close_stdin(self):
         user_input = dedent('''
             import os

Lib/test/test_tokenize.py

@@ -2427,6 +2427,30 @@ def test_stop_iteration_skips_encoded_readline_codec_lookup(self):
             (token.ENDMARKER, "", (1, 0), (1, 0), ""),
         )
 
+    def test_fstring_offsets_survive_buffer_reallocation(self):
+        for prefix in ("f", "t"):
+            for extra_tokens in (False, True):
+                with self.subTest(prefix=prefix, extra_tokens=extra_tokens):
+                    physical_lines = [
+                        prefix + '"""\n',
+                        "{(\n",
+                        " " * 9000 + "1\n",
+                        ")=:>{2}}\n",
+                        '"""\n',
+                    ]
+                    source = "".join(physical_lines)
+                    chunks = iter([
+                        "".join(physical_lines[:2]),
+                        "".join(physical_lines[2:4]),
+                        physical_lines[4],
+                        "",
+                    ])
+                    expected = self._get_tokens(
+                        source, extra_tokens=extra_tokens)
+                    tokens = list(tokenize._generate_tokens_from_c_tokenizer(
+                        chunks.__next__, extra_tokens=extra_tokens))
+                    self.assertEqual(tokens, expected)
+
     def test_extra_tokens_relaxes_lexer_errors(self):
         cases = [
             (

Modules/_testinternalcapi/tokenizer.c

@@ -195,6 +195,28 @@ test_tokenizer_source(PyObject *Py_UNUSED(module),
         goto error;
     }
 
+    _PyTok_SourceDiscard(&source);
+    if (check(_PyTok_SourceAppendLine(&source, "a\n", 2, 0) == 4,
+              "wrong retained source offset") < 0 ||
+            _PyTok_SourceLine(&source, 1, &line) < 0 ||
+            check(line.start == 4 && line.end == 6,
+                  "wrong retained source line") < 0 ||
+            _PyTok_SourceLocation(
+                &source, 4, _PYTOK_AFFINITY_LEFT, &loc) < 0 ||
+            check(loc.lineno == 1 && loc.byte_col == 0,
+                  "wrong retained source location") < 0) {
+        goto error;
+    }
+    view = _PyTok_SourceSpanView(
+        &source, _PyTok_SpanFromBounds(4, 5), &view_len);
+    if (check(view != NULL && view_len == 1 && view[0] == 'a',
+              "wrong retained source span") < 0 ||
+            check_system_error(_PyTok_SourceSpanView(
+                &source, _PyTok_SpanFromBounds(0, 1), &view_len) == NULL,
+                "accepted discarded source span") < 0) {
+        goto error;
+    }
+
     _PyTok_SourceClear(&source);
     Py_RETURN_NONE;
 
@@ -296,6 +318,88 @@ test_tokenizer_cursor(PyObject *Py_UNUSED(module),
     }
 #endif
 
+    _PyTok_Off base = source.len;
+    _PyTok_SourceDiscard(&source);
+    if (_PyTok_SourceAppendLine(&source, "ab\n", 3, 0) < 0 ||
+            _PyTok_SourceAppendLine(&source, "cd", 2, 0) < 0) {
+        goto error;
+    }
+    _PyTok_CursorInit(&cursor, &source);
+    if (_PyTok_CursorSetLine(&cursor, 1) < 0 ||
+            check(cursor.pos == base && _PyTok_CursorPeek(&cursor, 1) == 'b',
+                  "wrong retained cursor line") < 0 ||
+            _PyTok_CursorSetLine(&cursor, 2) < 0 ||
+            check(_PyTok_CursorAdvance(&cursor) == 'c',
+                  "wrong retained cursor byte") < 0 ||
+            _PyTok_CursorSetOffset(&cursor, base + 5) < 0 ||
+            check(cursor.lineno == 2 && _PyTok_CursorAdvance(&cursor) == EOF,
+                  "wrong retained cursor EOF") < 0) {
+        goto error;
+    }
+
+    _PyTok_SourceClear(&source);
+    Py_RETURN_NONE;
+
+error:
+    _PyTok_SourceClear(&source);
+    return NULL;
+}
+
+static PyObject *
+test_tokenizer_source_discard(PyObject *Py_UNUSED(module),
+                             PyObject *Py_UNUSED(args))
+{
+    _PyTok_SourceText source;
+    _PyTok_SourceInit(&source);
+    for (int i = 0; i < 260; i++) {
+        if (_PyTok_SourceAppendLine(&source, "x\n", 2, 1) < 0) {
+            goto error;
+        }
+    }
+    char *bytes = source.bytes;
+    _PyTok_Off capacity = source.cap;
+    _PyTok_SourceDiscard(&source);
+    if (check(source.base_offset == 520 && source.len == 0 &&
+                  source.nlines == 0 && source.bytes == bytes &&
+                  source.cap == capacity && source.bytes[0] == '\0',
+              "discard did not preserve source allocation") < 0) {
+        goto error;
+    }
+    for (int i = 0; i < 260; i++) {
+        if (check(_PyTok_SourceAppendLine(&source, "y\n", 2, 0) == 520 + 2 * i,
+                  "wrong source offset after discard") < 0 ||
+                check(!_PyTok_SourceLineIsImplicit(&source, i + 1),
+                      "discard preserved implicit newline flag") < 0) {
+            goto error;
+        }
+    }
+    if (check(source.bytes == bytes && source.cap == capacity,
+              "discarded allocation was not reused") < 0) {
+        goto error;
+    }
+    _PyTok_SourceDiscard(&source);
+    if (check(_PyTok_SourceAppendLine(&source, "tail", 4, 0) == 1040,
+              "wrong source offset after repeated discard") < 0) {
+        goto error;
+    }
+    _PyTok_SourceDiscard(&source);
+    if (check(_PyTok_SourceAppendLine(&source, "z\n", 2, 0) == 1044,
+              "cannot append after discarding unterminated line") < 0) {
+        goto error;
+    }
+    _PyTok_SourceDiscard(&source);
+    source.base_offset = PY_SSIZE_T_MAX - 1;
+    if (check(_PyTok_SourceAppendLine(&source, "z\n", 2, 0) < 0 &&
+                  PyErr_ExceptionMatches(PyExc_MemoryError),
+              "accepted overflowing logical source offset") < 0) {
+        goto error;
+    }
+    PyErr_Clear();
+    if (check(source.len == 0 && source.nlines == 0 &&
+                  source.base_offset == PY_SSIZE_T_MAX - 1,
+              "overflow changed retained source") < 0) {
+        goto error;
+    }
     _PyTok_SourceClear(&source);
     Py_RETURN_NONE;
 
@@ -307,6 +411,7 @@ test_tokenizer_cursor(PyObject *Py_UNUSED(module),
 static PyMethodDef test_methods[] = {
     {"test_tokenizer_source", test_tokenizer_source, METH_NOARGS},
     {"test_tokenizer_cursor", test_tokenizer_cursor, METH_NOARGS},
+    {"test_tokenizer_source_discard", test_tokenizer_source_discard, METH_NOARGS},
     {NULL},
 };
 

Parser/lexer/buffer.c

@@ -1,62 +1,46 @@
 #include "Python.h"
-#include "errcode.h"
-
+#include "buffer.h"
 #include "state.h"
 
-/* Traverse and remember all f-string buffers, in order to be able to restore
-   them after reallocating tok->buf */
 void
-_PyLexer_remember_fstring_buffers(struct tok_state *tok)
+_PyLexer_SaveBufferPointers(struct tok_state *tok, const char *base,
+                            _PyLexer_BufferPointers *pointers)
 {
-    int index;
-    tokenizer_mode *mode;
-
-    for (index = tok->tok_mode_stack_index; index >= 0; --index) {
-        mode = &(tok->tok_mode_stack[index]);
+    pointers->buf_from_base = tok->buf - base;
+    pointers->cur_from_buf = tok->cur - tok->buf;
+    pointers->inp_from_buf = tok->inp - tok->buf;
+    pointers->start_from_buf = tok->start == NULL
+        ? -1 : tok->start - tok->buf;
+    pointers->line_start_from_buf = tok->line_start == NULL
+        ? -1 : tok->line_start - tok->buf;
+    pointers->multi_line_start_from_buf = 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;
+        mode->multi_line_start_offset = mode->multi_line_start == NULL
+            ? -1 : mode->multi_line_start - tok->buf;
     }
 }
 
-/* Traverse and restore all f-string buffers after reallocating tok->buf */
 void
-_PyLexer_restore_fstring_buffers(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]);
-        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;
-    }
-}
-
-int
-_PyLexer_tok_reserve_buf(struct tok_state *tok, Py_ssize_t size)
+_PyLexer_RestoreBufferPointers(struct tok_state *tok, char *base,
+                               const _PyLexer_BufferPointers *pointers)
 {
-    Py_ssize_t cur = tok->cur - tok->buf;
-    Py_ssize_t oldsize = tok->inp - tok->buf;
-    Py_ssize_t newsize = oldsize + Py_MAX(size, oldsize >> 1);
-    if (newsize > tok->end - tok->buf) {
-        char *newbuf = tok->buf;
-        Py_ssize_t start = tok->start == NULL ? -1 : tok->start - tok->buf;
-        Py_ssize_t line_start = tok->start == NULL ? -1 : tok->line_start - tok->buf;
-        Py_ssize_t multi_line_start = tok->multi_line_start - tok->buf;
-        _PyLexer_remember_fstring_buffers(tok);
-        newbuf = (char *)PyMem_Realloc(newbuf, newsize);
-        if (newbuf == NULL) {
-            tok->done = E_NOMEM;
-            return 0;
-        }
-        tok->buf = newbuf;
-        tok->cur = tok->buf + cur;
-        tok->inp = tok->buf + oldsize;
-        tok->end = tok->buf + newsize;
-        tok->start = start < 0 ? NULL : tok->buf + start;
-        tok->line_start = line_start < 0 ? NULL : tok->buf + line_start;
-        tok->multi_line_start = multi_line_start < 0 ? NULL : tok->buf + multi_line_start;
-        _PyLexer_restore_fstring_buffers(tok);
+    tok->buf = base + pointers->buf_from_base;
+    tok->cur = tok->buf + pointers->cur_from_buf;
+    tok->inp = tok->buf + pointers->inp_from_buf;
+    tok->start = pointers->start_from_buf < 0
+        ? NULL : tok->buf + pointers->start_from_buf;
+    tok->line_start = pointers->line_start_from_buf < 0
+        ? NULL : tok->buf + pointers->line_start_from_buf;
+    tok->multi_line_start = pointers->multi_line_start_from_buf < 0
+        ? NULL : tok->buf + pointers->multi_line_start_from_buf;
+    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;
     }
-    return 1;
 }

Parser/lexer/buffer.h

@@ -3,8 +3,20 @@
 
 #include "pyport.h"
 
-void _PyLexer_remember_fstring_buffers(struct tok_state *tok);
-void _PyLexer_restore_fstring_buffers(struct tok_state *tok);
-int _PyLexer_tok_reserve_buf(struct tok_state *tok, Py_ssize_t size);
+struct tok_state;
+
+typedef struct {
+    Py_ssize_t buf_from_base;
+    Py_ssize_t cur_from_buf;
+    Py_ssize_t inp_from_buf;
+    Py_ssize_t start_from_buf;
+    Py_ssize_t line_start_from_buf;
+    Py_ssize_t multi_line_start_from_buf;
+} _PyLexer_BufferPointers;
+
+void _PyLexer_SaveBufferPointers(
+    struct tok_state *, const char *, _PyLexer_BufferPointers *);
+void _PyLexer_RestoreBufferPointers(
+    struct tok_state *, char *, const _PyLexer_BufferPointers *);
 
 #endif

Parser/lexer/lexer.c

@@ -12,8 +12,6 @@
 
 
 #define MAKE_TOKEN(token_type) _PyLexer_token_setup(tok, token, token_type, p_start, p_end)
-#define MAKE_TYPE_COMMENT_TOKEN(token_type, col_offset, end_col_offset) (\
-                _PyLexer_type_comment_token_setup(tok, token, token_type, col_offset, end_col_offset, p_start, p_end))
 
 /* Spaces in this constant are treated as "zero or more spaces or tabs" when
    tokenizing. */
@@ -360,21 +358,25 @@ _PyLexer_get_normal_mode(struct tok_state *tok, tokenizer_mode* current_tok, str
                     && !(tok->cur > ignore_end
                          && ((unsigned char)ignore_end[0] >= 128 || Py_ISALNUM(ignore_end[0]))));
 
+                int type = is_type_ignore ? TYPE_IGNORE : TYPE_COMMENT;
+                int start_col_offset = is_type_ignore
+                    ? ignore_end_col_offset : current_starting_col_offset;
+                p_end = tok->cur;
                 if (is_type_ignore) {
                     p_start = ignore_end;
-                    p_end = tok->cur;
 
                     /* If this type ignore is the only thing on the line, consume the newline also. */
                     if (blankline) {
                         tok_nextc(tok);
                         tok->atbol = 1;
                     }
-                    return MAKE_TYPE_COMMENT_TOKEN(TYPE_IGNORE, ignore_end_col_offset, tok->col_offset);
                 } else {
                     p_start = type_start;
-                    p_end = tok->cur;
-                    return MAKE_TYPE_COMMENT_TOKEN(TYPE_COMMENT, current_starting_col_offset, tok->col_offset);
                 }
+                _PyLexer_token_setup(tok, token, type, p_start, p_end);
+                token->start_loc = (_PyTok_Loc){tok->lineno, start_col_offset};
+                token->end_loc = (_PyTok_Loc){tok->lineno, tok->col_offset};
+                return type;
             }
         }
         if (tok->tok_extra_tokens) {

Parser/lexer/lexer.h

@@ -7,4 +7,25 @@ int _PyLexer_update_ftstring_expr(struct tok_state *tok, char cur);
 
 int _PyTokenizer_Get(struct tok_state *, struct token *);
 
+/* The view points into the current input window. The next
+   _PyTokenizer_Get() call may discard it. */
+static inline const char *
+_PyToken_TextView(const struct tok_state *tok, const struct token *token,
+                  Py_ssize_t *length)
+{
+    assert(length != NULL);
+    if (token->span.start < 0) {
+        assert(token->span.start == -1 && token->span.end == -1);
+        *length = 0;
+        return "";
+    }
+    assert(_PyTok_SpanIsValid(token->span));
+    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);
+}
+
 #endif