gh-153569: Use logical spans in tokenizer views and column calculations

Author: pablogsal

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 = [
             (

Parser/tokenizer/api.c

@@ -60,19 +60,14 @@ _PyToken_GetView(const struct tok_state *tok, const struct token *token,
                  _PyToken_View *view)
 {
     assert(view != NULL);
-    assert((token->span.start == -1 && token->span.end == -1) ||
-           _PyTok_SpanIsValid(token->span));
-    if (token->span.start >= 0) {
-        (void)_PyTok_SourcePointer(&tok->source, token->span.end);
-    }
-    view->text = token->span.start < 0
-        ? NULL : _PyTok_SourcePointer(&tok->source, token->span.start);
-    view->length = token->span.end - token->span.start;
-    view->end_line = _PyTok_SourcePointer(&tok->source, 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->text = _PyToken_TextView(tok, token, &view->length);
+    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;
 }

Parser/tokenizer/tokenizer.h

@@ -22,8 +22,8 @@ typedef struct {
     const char *text;
     Py_ssize_t length;
     const char *line;
-    Py_ssize_t line_length;
-    const char *end_line;
+    _PyTok_Span line_span;
+    _PyTok_Off end_line_start;
     int implicit_newline;
     int at_eof;
 } _PyToken_View;
@@ -73,8 +73,9 @@ _PyTokenizer_Info _PyTokenizer_GetInfo(const struct tok_state *);
 /* An absent token span has a nonnull empty text view. */
 const char *_PyToken_TextView(
     const struct tok_state *, const struct token *, Py_ssize_t *);
-/* Use the token from the most recent Get. text is NULL for an absent span;
-   line includes the token's complete physical line range. */
+/* Use the token from the most recent Get. An absent span has nonnull empty text.
+   line contains the bytes of line_span, the token's complete physical line
+   range. end_line_start is the logical offset of its final physical line. */
 void _PyToken_GetView(
     const struct tok_state *tok, const struct token *token,
     _PyToken_View *view);

Python/Python-tokenize.c

@@ -180,14 +180,13 @@ _get_col_offsets(tokenizeriterobject *it, const struct token *token,
                  Py_ssize_t *col_offset, Py_ssize_t *end_col_offset)
 {
     _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(it);
-    const char *token_start = view->text;
-    const char *token_end = token_start == NULL
-        ? NULL : token_start + view->length;
+    _PyTok_Off token_start = token->span.start;
+    _PyTok_Off token_end = token->span.end;
     Py_ssize_t lineno = token->start_loc.lineno;
     Py_ssize_t end_lineno = token->end_loc.lineno;
     Py_ssize_t byte_offset = -1;
-    if (token_start != NULL && token_start >= view->line) {
-        byte_offset = token_start - view->line;
+    if (token_start >= view->line_span.start) {
+        byte_offset = token_start - view->line_span.start;
         if (line_changed) {
             *col_offset = _PyPegen_byte_offset_to_character_offset_line(line, 0, byte_offset);
             if (*col_offset < 0) {
@@ -200,8 +199,8 @@ _get_col_offsets(tokenizeriterobject *it, const struct token *token,
         }
     }
 
-    if (token_end != NULL && token_end >= view->end_line) {
-        Py_ssize_t end_byte_offset = token_end - view->end_line;
+    if (token_end >= view->end_line_start) {
+        Py_ssize_t end_byte_offset = token_end - view->end_line_start;
         if (lineno == end_lineno) {
             // Avoid rescanning the prefix of a very long line.
             Py_ssize_t token_col_offset = _PyPegen_byte_offset_to_character_offset_line(line, byte_offset, end_byte_offset);
@@ -213,7 +212,8 @@ _get_col_offsets(tokenizeriterobject *it, const struct token *token,
         }
         else {
             *end_col_offset = _PyPegen_byte_offset_to_character_offset_line(
-                line, view->end_line - view->line, token_end - view->line);
+                line, view->end_line_start - view->line_span.start,
+                token_end - view->line_span.start);
             if (*end_col_offset < 0) {
                 return -1;
             }
@@ -251,15 +251,7 @@ tokenizeriter_next(PyObject *op)
     }
     _PyToken_View view;
     _PyToken_GetView(it->tok, &token, &view);
-    const char *token_start = view.text;
-    PyObject *str;
-    if (token.span.start < 0) {
-        assert(token.span.start == -1 && token.span.end == -1);
-        str = Py_GetConstant(Py_CONSTANT_EMPTY_STR);
-    }
-    else {
-        str = PyUnicode_FromStringAndSize(token_start, view.length);
-    }
+    PyObject *str = PyUnicode_FromStringAndSize(view.text, view.length);
     if (str == NULL) {
         goto exit;
     }
@@ -271,7 +263,7 @@ tokenizeriter_next(PyObject *op)
     if (it->extra_tokens && is_trailing_token) {
         line = Py_GetConstant(Py_CONSTANT_EMPTY_STR);
     } else {
-        Py_ssize_t size = view.line_length;
+        Py_ssize_t size = view.line_span.end - view.line_span.start;
         if (size >= 1 && view.implicit_newline) {
             size -= 1;
         }
@@ -307,8 +299,7 @@ tokenizeriter_next(PyObject *op)
         else if (type == NEWLINE) {
             if (!view.implicit_newline) {
                 Py_DECREF(str);
-                assert(token_start != NULL);
-                if (token_start[0] == '\r') {
+                if (view.text[0] == '\r') {
                     str = PyUnicode_FromString("\r\n");
                 } else {
                     str = PyUnicode_FromString("\n");