gh-153569: Simplify tokenizer source views and reader ownership

Author: pablogsal

Lib/test/test_capi/test_tokenizer.py

@@ -12,6 +12,9 @@ def test_source(self):
     def test_source_discard(self):
         _testinternalcapi.test_tokenizer_source_discard()
 
+    def test_source_views(self):
+        _testinternalcapi.test_tokenizer_source_views()
+
 
 if __name__ == "__main__":
     unittest.main()

Lib/test/test_tokenize.py

@@ -2498,6 +2498,34 @@ def test_fstring_offsets_survive_buffer_reallocation(self):
                         chunks.__next__, extra_tokens=extra_tokens))
                     self.assertEqual(tokens, expected)
 
+    def test_multiline_unicode_columns_after_source_discard(self):
+        # The first line is discarded before the multiline token is read.
+        # Its saved offsets must remain valid when the source buffer grows.
+        padding = "é" * 9000
+        source = f"pass\né = '''{padding}\n漢''' + ö\nß = 1\n"
+        expected = [
+            (token.NAME, "pass", (1, 0), (1, 4)),
+            (token.NAME, "é", (2, 0), (2, 1)),
+            (token.STRING, f"'''{padding}\n漢'''", (2, 4), (3, 4)),
+            (token.NAME, "ö", (3, 7), (3, 8)),
+            (token.NAME, "ß", (4, 0), (4, 1)),
+        ]
+        for extra_tokens in (False, True):
+            for encoding in (None, "utf-8"):
+                with self.subTest(extra_tokens=extra_tokens, encoding=encoding):
+                    stream = (StringIO(source) if encoding is None
+                              else BytesIO(source.encode(encoding)))
+                    tokens = list(tokenize._generate_tokens_from_c_tokenizer(
+                        stream.readline, extra_tokens=extra_tokens,
+                        encoding=encoding))
+                    self.assertEqual(
+                        [t[:4] for t in tokens
+                         if t.type in (token.NAME, token.STRING)],
+                        expected,
+                    )
+                    string = next(t for t in tokens if t.type == token.STRING)
+                    self.assertEqual(string.line, f"é = '''{padding}\n漢''' + ö\n")
+
     def test_extra_tokens_relaxes_lexer_errors(self):
         cases = [
             (

Modules/_testinternalcapi/tokenizer.c

@@ -103,6 +103,52 @@ test_tokenizer_source(PyObject *Py_UNUSED(module),
     return NULL;
 }
 
+static PyObject *
+test_tokenizer_source_views(PyObject *Py_UNUSED(module),
+                           PyObject *Py_UNUSED(args))
+{
+    const char line[] = "\xce\xb2\n";
+    _PyTok_SourceText source;
+    _PyTok_SourceInit(&source);
+    if (_PyTok_SourceAppendLine(&source, line, sizeof(line) - 1) < 0) {
+        goto error;
+    }
+    _PyTok_SourceDiscard(&source);
+    _PyTok_Off start = _PyTok_SourceAppendLine(&source, line, sizeof(line) - 1);
+    if (start < 0) {
+        goto error;
+    }
+    _PyTok_Span span = {start, start + (Py_ssize_t)sizeof(line) - 1};
+    // Grow storage after saving a span with a nonzero logical base.
+    Py_ssize_t capacity = source.cap;
+    while (source.cap == capacity) {
+        if (_PyTok_SourceAppendLine(&source, line, sizeof(line) - 1) < 0) {
+            goto error;
+        }
+    }
+    Py_ssize_t len;
+    const char *view = _PyTok_SourceSpanView(&source, span, &len);
+    if (check(len == (Py_ssize_t)sizeof(line) - 1 && memcmp(view, line, len) == 0,
+              "source span changed after growth") < 0 ||
+            check(_PyTok_SourceOffset(&source, view) == start &&
+                  _PyTok_SourcePointer(&source, span.end) == view + len,
+                  "source view lost its logical offset") < 0) {
+        goto error;
+    }
+    _PyTok_Off end = source.base_offset + source.len;
+    view = _PyTok_SourceSpanView(&source, (_PyTok_Span){end, end}, &len);
+    if (check(len == 0 && view == source.bytes + source.len,
+              "wrong empty view at source end") < 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))
@@ -180,6 +226,7 @@ test_tokenizer_source_discard(PyObject *Py_UNUSED(module),
 
 static PyMethodDef test_methods[] = {
     {"test_tokenizer_source", test_tokenizer_source, METH_NOARGS},
+    {"test_tokenizer_source_views", test_tokenizer_source_views, METH_NOARGS},
     {"test_tokenizer_source_discard", test_tokenizer_source_discard, METH_NOARGS},
     {NULL},
 };

Parser/lexer/lexer.c

@@ -29,8 +29,9 @@ _PyLexer_refill(struct tok_state *tok)
 #if defined(Py_DEBUG)
     if (tok->debug) {
         fprintf(stderr, "line[%d] = ", tok->lineno);
-        _PyTokenizer_print_escape(stderr, _PyLexer_BufferPointer(tok, tok->cur),
-                                  tok->inp - tok->cur);
+        _PyTokenizer_print_escape(
+            stderr, _PyTok_SourcePointer(&tok->source, tok->cur),
+            tok->inp - tok->cur);
         fprintf(stderr, "  tok->done = %d\n", tok->done);
     }
 #endif
@@ -39,7 +40,7 @@ _PyLexer_refill(struct tok_state *tok)
         return 0;
     }
     tok->line_start = tok->cur;
-    if (contains_null_bytes(_PyLexer_BufferPointer(tok, tok->line_start),
+    if (contains_null_bytes(_PyTok_SourcePointer(&tok->source, tok->line_start),
                             tok->inp - tok->line_start)) {
         _PyTokenizer_syntaxerror(tok, "source code cannot contain null bytes");
         tok->cur = tok->inp;
@@ -56,7 +57,8 @@ _PyLexer_backup(struct tok_state *tok, int c)
         if (--tok->cur < tok->buf_offset) {
             Py_FatalError("tokenizer beginning of buffer");
         }
-        if ((int)(unsigned char)*_PyLexer_BufferPointer(tok, tok->cur) != Py_CHARMASK(c)) {
+        const char *cur = _PyTok_SourcePointer(&tok->source, tok->cur);
+        if ((int)(unsigned char)*cur != Py_CHARMASK(c)) {
             Py_FatalError("tok_backup: wrong character");
         }
     }
@@ -74,7 +76,8 @@ verify_identifier(struct tok_state *tok)
     PyObject *s;
     if (tok_failed(tok))
         return 0;
-    s = PyUnicode_DecodeUTF8(_PyLexer_BufferPointer(tok, tok->start), tok->cur - tok->start, NULL);
+    s = PyUnicode_DecodeUTF8(_PyTok_SourcePointer(&tok->source, tok->start),
+                             tok->cur - tok->start, NULL);
     if (s == NULL) {
         if (PyErr_ExceptionMatches(PyExc_UnicodeDecodeError)) {
             tok->done = E_DECODE;
@@ -105,13 +108,13 @@ verify_identifier(struct tok_state *tok)
         Py_DECREF(s);
         if (Py_UNICODE_ISPRINTABLE(ch)) {
             _PyTokenizer_syntaxerror_at(
-                tok, _PyLexer_BufferPointer(tok, tok->line_start),
+                tok, _PyTok_SourcePointer(&tok->source, tok->line_start),
                 error_cursor - tok->line_start, tok->lineno, -1, -1,
                 "invalid character '%c' (U+%04X)", ch, ch);
         }
         else {
             _PyTokenizer_syntaxerror_at(
-                tok, _PyLexer_BufferPointer(tok, tok->line_start),
+                tok, _PyTok_SourcePointer(&tok->source, tok->line_start),
                 error_cursor - tok->line_start, tok->lineno, -1, -1,
                 "invalid non-printable character U+%04X", ch);
         }
@@ -193,14 +196,14 @@ _PyLexer_get_normal(struct tok_state *tok, ftstring_state *current, struct token
         }
 
         if (tok->tok_extra_tokens) {
-            p = _PyLexer_BufferPointer(tok, tok->start);
+            p = _PyTok_SourcePointer(&tok->source, tok->start);
         }
 
         if (tok->type_comments) {
-            p = _PyLexer_BufferPointer(tok, tok->start);
+            p = _PyTok_SourcePointer(&tok->source, tok->start);
             current_starting_col_offset = tok->start_loc.byte_col;
             prefix = type_comment_prefix;
-            while (*prefix && p < _PyLexer_BufferPointer(tok, tok->cur)) {
+            while (*prefix && p < _PyTok_SourcePointer(&tok->source, tok->cur)) {
                 if (*prefix == ' ') {
                     while (*p == ' ' || *p == '\t') {
                         p++;
@@ -229,24 +232,25 @@ _PyLexer_get_normal(struct tok_state *tok, ftstring_state *current, struct token
                 /* A TYPE_IGNORE is "type: ignore" followed by the end of the token
                  * or anything ASCII and non-alphanumeric. */
                 is_type_ignore = (
-                    _PyLexer_BufferPointer(tok, tok->cur) >= ignore_end && memcmp(p, "ignore", 6) == 0
-                    && !(_PyLexer_BufferPointer(tok, tok->cur) > ignore_end
+                    _PyTok_SourcePointer(&tok->source, tok->cur) >= ignore_end
+                    && memcmp(p, "ignore", 6) == 0
+                    && !(_PyTok_SourcePointer(&tok->source, 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 = _PyLexer_BufferOffset(tok, ignore_end);
+                    p_start = _PyTok_SourceOffset(&tok->source, ignore_end);
 
                     /* If this type ignore is the only thing on the line, consume the newline also. */
                     if (blankline) {
                         tok_nextc(tok);
                         tok->layout.at_bol = 1;
                     }
                 } else {
-                    p_start = _PyLexer_BufferOffset(tok, type_start);
+                    p_start = _PyTok_SourceOffset(&tok->source, type_start);
                 }
                 _PyLexer_token_setup(tok, token, type, p_start, p_end);
                 token->start_loc = (_PyTok_Loc){tok->lineno, start_col_offset};
@@ -257,7 +261,7 @@ _PyLexer_get_normal(struct tok_state *tok, ftstring_state *current, struct token
         }
         if (tok->tok_extra_tokens) {
             tok_backup(tok, c);  /* don't eat the newline or EOF */
-            p_start = _PyLexer_BufferOffset(tok, p);
+            p_start = _PyTok_SourceOffset(&tok->source, p);
             p_end = tok->cur;
             tok->layout.comment_newline = blankline;
             return MAKE_TOKEN(COMMENT);

Parser/lexer/state.c

@@ -58,9 +58,6 @@ _PyLexer_PopFTString(struct tok_state *tok)
 void
 _PyTokenizer_Free(struct tok_state *tok)
 {
-    if (tok->encoding != NULL) {
-        PyMem_Free(tok->encoding);
-    }
     Py_XDECREF(tok->filename);
     Py_XDECREF(tok->module);
     _PyTok_ReaderFree(tok);

Parser/lexer/state.h

@@ -80,7 +80,6 @@ struct tok_state {
     _PyTok_SourceText source;
     int done;           /* E_OK normally, E_EOF at EOF, otherwise error code */
     /* NB If done != E_OK, cur must be == inp!!! */
-    FILE *fp;           /* Rest of input; NULL if tokenizing a string */
     lexer_layout_state layout;
     int lineno;         /* Current line number */
     _PyTok_Loc start_loc;
@@ -92,9 +91,6 @@ struct tok_state {
     int parencolstack[MAXLEVEL];
     PyObject *filename;
     PyObject *module;
-    /* Stuff for PEP 0263 */
-    char *encoding;         /* Source encoding. */
-
     struct _PyTok_Reader *reader;
 
     int type_comments;      /* Whether to look for type comments */
@@ -135,33 +131,6 @@ _PyLexer_FTStringBracketDepth(const struct tok_state *tok,
     return tok->level - state->paren_level;
 }
 
-static inline _PyTok_Off
-_PyLexer_BufferOffset(const struct tok_state *tok, const char *position)
-{
-    const char *base = _PyTok_SourceData(&tok->source);
-    assert(position >= base && position <= base + tok->source.len);
-    return tok->source.base_offset + (position - base);
-}
-
-static inline const char *
-_PyLexer_BufferPointer(const struct tok_state *tok, _PyTok_Off offset)
-{
-    assert(offset >= tok->source.base_offset);
-    assert(offset - tok->source.base_offset <= tok->source.len);
-    return _PyTok_SourceData(&tok->source) + (offset - tok->source.base_offset);
-}
-
-static inline const char *
-_PyLexer_BufferSpanView(const struct tok_state *tok, _PyTok_Span span,
-                        Py_ssize_t *length)
-{
-    assert(length != NULL);
-    assert(_PyTok_SpanIsValid(span));
-    *length = span.end - span.start;
-    (void)_PyLexer_BufferPointer(tok, span.end);
-    return _PyLexer_BufferPointer(tok, span.start);
-}
-
 static inline int
 _PyLexer_ByteColumn(const struct tok_state *tok)
 {

Parser/lexer/string.c

@@ -77,8 +77,8 @@ finish_ftstring_expr(struct tok_state *tok, ftstring_state *state,
         return 0;
     }
     Py_ssize_t expr_len;
-    const char *expr = _PyLexer_BufferSpanView(
-        tok, state->expr_span, &expr_len);
+    const char *expr = _PyTok_SourceSpanView(
+        &tok->source, state->expr_span, &expr_len);
     tokenizer_comments *comments = state->comments;
     PyObject *res;
     if (comments != NULL && comments->count > 0) {
@@ -339,7 +339,8 @@ _PyLexer_scan_string(struct tok_state *tok, struct token *token, int c)
             }
             int end_lineno = tok->lineno;
             _PyTok_Loc location = tok->start_loc;
-            const char *line = _PyLexer_BufferPointer(tok, tok->start) - location.byte_col;
+            const char *line = _PyTok_SourcePointer(
+                &tok->source, tok->start - location.byte_col);
             Py_ssize_t cursor_offset = (Py_ssize_t)location.byte_col + 1;
 
             const ftstring_state *state = _PyLexer_CurrentFTString(tok);
@@ -460,7 +461,8 @@ _PyLexer_get_ftstring(struct tok_state *tok, ftstring_state *current, struct tok
 
             int end_lineno = tok->lineno;
             _PyTok_Loc location = current->start_loc;
-            const char *line = _PyLexer_BufferPointer(tok, current->start) - location.byte_col;
+            const char *line = _PyTok_SourcePointer(
+                &tok->source, current->start - location.byte_col);
             Py_ssize_t cursor_offset = (Py_ssize_t)location.byte_col + 1;
 
             if (quote_size == 3) {

Parser/tokenizer/api.c

@@ -4,6 +4,7 @@
 
 #include "tokenizer.h"
 #include "reader.h"
+#include "reader_internal.h"
 #include "../lexer/state.h"
 
 _PyTokenizer_Info
@@ -21,10 +22,10 @@ _PyTokenizer_GetInfo(const struct tok_state *tok)
         .delimiter_loc = {-1, -1},
         .in_formatted_string = tok->ftstring_depth != 0,
         .is_interactive = _PyTok_ReaderIsInteractive(tok),
-        .is_file = tok->fp != NULL && tok->fp != stdin,
+        .is_file = tok->reader->fp != NULL && tok->reader->fp != stdin,
         .filename = tok->filename,
         .module = tok->module,
-        .encoding = tok->encoding,
+        .encoding = tok->reader->encoding,
     };
     if (tok->level > 0) {
         int level = tok->level - 1;
@@ -45,14 +46,14 @@ _PyToken_TextView(const struct tok_state *tok, const struct token *token,
         *length = 0;
         return "";
     }
-    return _PyLexer_BufferSpanView(tok, token->span, length);
+    return _PyTok_SourceSpanView(&tok->source, token->span, length);
 }
 
 const char *
 _PyTokenizer_SpanView(const struct tok_state *tok, _PyTok_Span span,
                       Py_ssize_t *length)
 {
-    return _PyLexer_BufferSpanView(tok, span, length);
+    return _PyTok_SourceSpanView(&tok->source, span, length);
 }
 
 void
@@ -63,16 +64,18 @@ _PyToken_GetView(const struct tok_state *tok, const struct token *token,
     assert((token->span.start == -1 && token->span.end == -1) ||
            _PyTok_SpanIsValid(token->span));
     if (token->span.start >= 0) {
-        (void)_PyLexer_BufferPointer(tok, token->span.end);
+        (void)_PyTok_SourcePointer(&tok->source, token->span.end);
     }
     view->text = token->span.start < 0
-        ? NULL : _PyLexer_BufferPointer(tok, token->span.start);
+        ? NULL : _PyTok_SourcePointer(&tok->source, token->span.start);
     view->length = token->span.end - token->span.start;
-    view->end_line = _PyLexer_BufferPointer(tok, tok->line_start);
-    view->line = ISSTRINGLIT(token->type)
-        ? view->text - token->start_loc.byte_col : view->end_line;
-    view->line_length = tok->inp - tok->line_start +
-        (view->end_line - view->line);
+    view->end_line_start = tok->line_start;
+    view->line_span = (_PyTok_Span){
+        ISSTRINGLIT(token->type)
+            ? token->span.start - token->start_loc.byte_col : tok->line_start,
+        tok->inp,
+    };
+    view->line = _PyTok_SourcePointer(&tok->source, view->line_span.start);
     view->implicit_newline = tok->implicit_newline;
     view->at_eof = tok->done == E_EOF;
 }
@@ -111,7 +114,7 @@ _PyTokenizer_ImplyDedents(struct tok_state *tok)
 int
 _PyTokenizer_HasTrailingStatement(const struct tok_state *tok)
 {
-    const char *cur = _PyLexer_BufferPointer(tok, tok->cur);
+    const char *cur = _PyTok_SourcePointer(&tok->source, tok->cur);
     char c = *cur;
     for (;;) {
         while (c == ' ' || c == '\t' || c == '\n' || c == '\014') {