gh-153569: tighten tokenizer iterator cleanup

Author: pablogsal

Lib/test/test_tokenize.py

@@ -2282,6 +2282,39 @@ def showwarning(*args, **kwargs):
                 ):
                     next(iterator)
 
+    def test_exhausted_iterator_releases_readline(self):
+        class Readline:
+            def __init__(self):
+                self.lines = iter(["x\n", ""])
+
+            def __call__(self):
+                return next(self.lines)
+
+        readline = Readline()
+        readline_ref = weakref.ref(readline)
+        iterator = _tokenize.TokenizerIter(readline, extra_tokens=False)
+        del readline
+
+        for _ in iterator:
+            pass
+        support.gc_collect()
+
+        self.assertIsNone(readline_ref())
+        with self.assertRaises(StopIteration):
+            next(iterator)
+
+    def test_lone_carriage_return(self):
+        self.assertEqual(
+            list(tokenize.generate_tokens(StringIO("x\r").readline)),
+            [
+                tokenize.TokenInfo(token.NAME, "x", (1, 0), (1, 1), "x\r"),
+                tokenize.TokenInfo(
+                    token.NEWLINE, "\r", (1, 1), (1, 3), "x\r"
+                ),
+                tokenize.TokenInfo(token.ENDMARKER, "", (2, 0), (2, 0), ""),
+            ],
+        )
+
     def check_tokenize(self, s, expected):
         # Format the tokens in s in a table format.
         # The ENDMARKER and final NEWLINE are omitted.

Parser/tokenizer/api.c

@@ -117,8 +117,8 @@ _PyTok_TokenTextLocations(const PyTokenizer *tok, const _PyTok_Token *token,
         *end = token->end;
         return 0;
     }
-    if (_PyTok_SourceLocation(
-            &tok->source, token->span.start,
+    if (_PyTok_CursorLocation(
+            &tok->cursor, token->span.start,
             _PYTOK_AFFINITY_RIGHT, start) < 0) {
         return -1;
     }
@@ -127,8 +127,8 @@ _PyTok_TokenTextLocations(const PyTokenizer *tok, const _PyTok_Token *token,
         token->span.start == token->span.end
         ? _PYTOK_AFFINITY_RIGHT
         : _PYTOK_AFFINITY_LEFT;
-    return _PyTok_SourceLocation(
-        &tok->source, token->span.end, end_affinity, end);
+    return _PyTok_CursorLocation(
+        &tok->cursor, token->span.end, end_affinity, end);
 }
 
 static int

Parser/tokenizer/reader.c

@@ -727,19 +727,6 @@ tokenizer_from_string(const char *input, int utf8_only, int exec_input,
     return tok;
 }
 
-static struct _PyTokenizer *
-tokenizer_from_plain_string(const char *input, int exec_input,
-                            int preserve_crlf)
-{
-    return tokenizer_from_string(input, 0, exec_input, preserve_crlf);
-}
-
-static struct _PyTokenizer *
-tokenizer_from_utf8(const char *input, int exec_input, int preserve_crlf)
-{
-    return tokenizer_from_string(input, 1, exec_input, preserve_crlf);
-}
-
 static struct _PyTokenizer *
 tokenizer_from_readline(PyObject *readline, const char *encoding)
 {
@@ -788,12 +775,12 @@ _PyTok_ReaderNew(const _PyTok_Config *config)
 {
     switch (config->kind) {
         case _PYTOK_SOURCE_STRING:
-            return tokenizer_from_plain_string(
-                config->source.string, config->exec_input,
+            return tokenizer_from_string(
+                config->source.string, 0, config->exec_input,
                 config->preserve_crlf);
         case _PYTOK_SOURCE_UTF8:
-            return tokenizer_from_utf8(
-                config->source.string, config->exec_input,
+            return tokenizer_from_string(
+                config->source.string, 1, config->exec_input,
                 config->preserve_crlf);
         case _PYTOK_SOURCE_FILE:
             return tokenizer_from_file(

Python/Python-tokenize.c

@@ -211,6 +211,8 @@ tokenizeriter_next(PyObject *op)
 {
     tokenizeriterobject *it = (tokenizeriterobject*)op;
     PyObject* result = NULL;
+    PyTokenizer *tok_to_free = NULL;
+    PyObject *line_to_clear = NULL;
 
     Py_BEGIN_CRITICAL_SECTION(it);
 
@@ -279,17 +281,12 @@ tokenizeriter_next(PyObject *op)
     }
 
     if (_PyTok_ExtraTokens(it->tok)) {
-        // Necessary adjustments to match the original Python tokenize
-        // implementation
         if (type > DEDENT && type < OP) {
             type = OP;
         }
         else if (type == NEWLINE) {
-            Py_DECREF(str);
-            if (token.flags & _PYTOK_IMPLICIT_NL) {
-                str = Py_GetConstant(Py_CONSTANT_EMPTY_STR);
-            }
-            else {
+            if (!(token.flags & _PYTOK_IMPLICIT_NL)) {
+                Py_DECREF(str);
                 if (token_len > 0 && token_view[0] == '\r') {
                     str = PyUnicode_FromString("\r\n");
                 }
@@ -318,9 +315,15 @@ tokenizeriter_next(PyObject *op)
     _PyTok_TokenClear(&token);
     if (type == ENDMARKER) {
         it->done = 1;
+        tok_to_free = it->tok;
+        it->tok = NULL;
+        line_to_clear = it->last_line;
+        it->last_line = NULL;
     }
 
     Py_END_CRITICAL_SECTION();
+    _PyTok_Free(tok_to_free);
+    Py_XDECREF(line_to_clear);
     return result;
 }