gh-159672: Preserve error state during keyword replay

Author: pablogsal

Lib/test/test_traceback.py

@@ -1833,6 +1833,11 @@ class TestKeywordTypoSuggestions(unittest.TestCase):
         ("void f():", "def"),
     ]
 
+    def get_syntax_error(self, source, mode="exec"):
+        with self.assertRaises(SyntaxError) as caught:
+            compile(source, "<test>", mode)
+        return caught.exception
+
     def test_keyword_suggestions_from_file(self):
         with tempfile.TemporaryDirectory() as script_dir:
             for i, (code, expected_kw) in enumerate(self.TYPO_CASES):
@@ -1856,9 +1861,8 @@ def test_keyword_suggestions_from_single_input(self):
         for source, keyword in self.TYPO_CASES:
             for ending in ('', '\n'):
                 with self.subTest(source=source, ending=ending):
-                    with self.assertRaises(SyntaxError) as caught:
-                        compile(source + ending, "<test>", "single")
-                    self.assertEqual(caught.exception._metadata[0], keyword)
+                    error = self.get_syntax_error(source + ending, "single")
+                    self.assertEqual(error._metadata[0], keyword)
 
     def test_no_keyword_suggestion_for_comma_errors(self):
         # When the parser identifies a missing comma, don't suggest
@@ -1897,9 +1901,7 @@ def test_parser_suggestions_need_no_formatting_work(self):
         ]
         for source, keyword, lineno, offset in cases:
             with self.subTest(source=source):
-                with self.assertRaises(SyntaxError) as caught:
-                    compile(source, "<test>", "exec")
-                error = caught.exception
+                error = self.get_syntax_error(source)
                 original = error.args
                 with (unittest.mock.patch("builtins.compile") as compile_mock,
                       unittest.mock.patch("builtins.open") as open_mock):
@@ -1912,28 +1914,88 @@ def test_parser_suggestions_need_no_formatting_work(self):
                 self.assertEqual((exc.lineno, exc.offset), (str(lineno), offset))
                 self.assertEqual(error.args, original)
 
+    @cpython_only
+    def test_public_fields_override_keyword_metadata(self):
+        changes = dict(filename="changed.py", msg="custom message", text="<redacted>\n",
+                       lineno=101, offset=2, end_lineno=102, end_offset=3)
+        for field, value in changes.items():
+            with self.subTest(field=field):
+                error = self.get_syntax_error("impor math")
+                setattr(error, field, value)
+                error.args = (error.msg, (error.filename, error.lineno, error.offset,
+                                         error.text, error.end_lineno, error.end_offset))
+                formatted = traceback.TracebackException.from_exception(error)
+                expected = str(value) if field in ("lineno", "end_lineno") else value
+                self.assertEqual(getattr(formatted, field), expected)
+                self.assertNotIn("Did you mean", ''.join(formatted.format_exception_only()))
+
+    @cpython_only
+    def test_keyword_metadata_avoids_custom_methods(self):
+        class CustomString(str):
+            def __eq__(self, other):
+                raise AssertionError("metadata called a custom method")
+
+            def __str__(self):
+                raise AssertionError("metadata called a custom method")
+
+        error = self.get_syntax_error("impor math")
+        suggestion, location, original = error._metadata
+        for metadata in ((), (CustomString(suggestion), location, original),
+                         (suggestion, (CustomString(location[0]), *location[1:]), original),
+                         (suggestion, location, (CustomString(original[0]), *original[1:]))):
+            with self.subTest(metadata_type=type(metadata[0]) if metadata else None):
+                error._metadata = metadata
+                formatted = ''.join(traceback.format_exception_only(error))
+                self.assertNotIn("Did you mean", formatted)
+
+    def test_keyword_hint_in_interactive_interpreter(self):
+        import code
+        with captured_output("stderr") as output:
+            code.InteractiveInterpreter().runsource("impor math", filename="user.py")
+        self.assertIn("Did you mean 'import'", output.getvalue())
+
+    @requires_subprocess()
+    def test_keyword_hint_does_not_reopen_source(self):
+        with tempfile.TemporaryDirectory() as script_dir:
+            make_script(script_dir, "sitecustomize", textwrap.dedent('''
+                import sys
+                def audit(event, args):
+                    if event == "open" and isinstance(args[0], str) and args[0].endswith("script.py"):
+                        print("source-open")
+                sys.addaudithook(audit)
+            '''))
+            counts = []
+            cases = [("= math", None), ("impor math", "import"),
+                     ("result = (one\n          adn two)", "and")]
+            for source, keyword in cases:
+                script = make_script(script_dir, "script", source)
+                _, output, error = assert_python_failure(script, PYTHONPATH=script_dir)
+                counts.append(output.count(b"source-open"))
+                if keyword is not None:
+                    self.assertIn(f"Did you mean '{keyword}'".encode(), error)
+                if keyword == "and":
+                    self.assertIn(b"line 1\n    result = (one\n", error)
+            self.assertGreater(counts[0], 0)
+            self.assertEqual(counts, [counts[0]] * len(cases))
+
     @cpython_only
     def test_keyword_trials_do_not_repeat_warnings(self):
         import warnings
         for prefix in ("value = '\\q'\n", "from . lazy import x\n"):
             with self.subTest(prefix=prefix), warnings.catch_warnings(record=True) as caught:
                 warnings.simplefilter("always")
-                with self.assertRaises(SyntaxError) as error:
-                    compile(prefix + "impor math", "<test>", "exec")
+                error = self.get_syntax_error(prefix + "impor math")
                 self.assertEqual(len(caught), 1)
-                self.assertEqual(error.exception._metadata[0], "import")
+                self.assertEqual(error._metadata[0], "import")
 
     @cpython_only
     def test_keyword_replay_requires_complete_valid_input(self):
         for source in ("dell 1", "impor math\nother =", "whille True:\npass",
                        "switch +:\n case:", "function f(=):\n pass", "(mat x)",
                        "impor math\n" + "pass\n" * 20000):
             with self.subTest(source=source):
-                with self.assertRaises(SyntaxError) as caught:
-                    compile(source, "<test>", "exec")
-                # Action errors and a cache which stopped before EOF must not
-                # count as successful repairs.
-                self.assertIsNone(caught.exception._metadata)
+                error = self.get_syntax_error(source)
+                self.assertIsNone(error._metadata)
 
     @cpython_only
     def test_keyword_replay_respects_feature_version(self):
@@ -1949,9 +2011,8 @@ def test_keyword_trials_reject_invalid_literals(self):
             for mode in ("exec", "eval", "single"):
                 source = "npt " + literal
                 with self.subTest(source=source, mode=mode):
-                    with self.assertRaises(SyntaxError) as caught:
-                        compile(source, "<test>", mode)
-                    self.assertIsNone(caught.exception._metadata)
+                    error = self.get_syntax_error(source, mode)
+                    self.assertIsNone(error._metadata)
 
     @cpython_only
     @requires_subprocess()
@@ -2008,17 +2069,15 @@ def test_keyword_trials_numeric_limit(self):
             with self.subTest(digits=digits):
                 source = "x = " + "1" * digits
                 compile(source, "<test>", "exec")
-                with self.assertRaises(SyntaxError) as caught:
-                    compile("impor math\n" + source, "<test>", "exec")
-                self.assertEqual(caught.exception._metadata is not None, expected)
+                error = self.get_syntax_error("impor math\n" + source)
+                self.assertEqual(error._metadata is not None, expected)
 
     @cpython_only
     def test_keyword_trials_do_not_call_normalizer(self):
         with unittest.mock.patch("unicodedata.normalize") as normalize:
-            with self.assertRaises(SyntaxError) as caught:
-                compile("whille True:\n    café = 1", "<test>", "exec")
+            error = self.get_syntax_error("whille True:\n    café = 1")
         normalize.assert_not_called()
-        self.assertIsNone(caught.exception._metadata)
+        self.assertIsNone(error._metadata)
 
     @cpython_only
     @requires_subprocess()
@@ -2051,10 +2110,9 @@ def test_keyword_replay_source_limit(self):
                 literal = "value = " + repr("a" * size) + "\n"
                 for source in (literal + "impor math", "impor math\n" + literal):
                     with self.subTest(size=size, typo_first=source.startswith("impor")):
-                        with self.assertRaises(SyntaxError) as caught:
-                            compile(source, "<test>", "exec")
+                        error = self.get_syntax_error(source)
                         expected = size < 256 * 1024
-                        self.assertEqual(caught.exception._metadata is not None, expected)
+                        self.assertEqual(error._metadata is not None, expected)
                         script = make_script(script_dir, "large_keyword_typo", source)
                         _, _, stderr = assert_python_failure(script)
                         self.assertEqual(b"Did you mean 'import'" in stderr, expected)
@@ -2066,10 +2124,9 @@ def test_keyword_replay_work_limit(self):
         # Several candidates must share the budget, including failed passes.
         for count, expected in ((100, "lambda"), (2000, None)):
             source = "x = 1\n" * count + "f(" + "fron, " * 8 + "lamda x: x)"
-            with self.subTest(count=count), self.assertRaises(SyntaxError) as caught:
-                compile(source, "<test>", "exec")
-            metadata = caught.exception._metadata
-            self.assertEqual(metadata[0] if metadata else None, expected)
+            with self.subTest(count=count):
+                metadata = self.get_syntax_error(source)._metadata
+                self.assertEqual(metadata[0] if metadata else None, expected)
 
 @requires_debug_ranges()
 @force_not_colorized_test_class

Lib/traceback.py

@@ -1086,6 +1086,27 @@ def emit(self, text_gen, margin_char=None):
                 yield textwrap.indent(text, indent_str, lambda line: True)
 
 
+def _syntax_error_metadata_matches(metadata, fields):
+    if type(metadata) is not tuple or len(metadata) != 3:
+        return False
+    suggestion, location, original = metadata
+    if type(suggestion) is not str or type(original) is not tuple:
+        return False
+    if len(original) != len(fields):
+        return False
+    if location is not None:
+        if type(location) is not tuple or len(location) != 5:
+            return False
+        if not all(type(value) is expected for value, expected in
+                   zip(location, (str, int, int, int, int))):
+            return False
+    # Public fields can change independently of args. Exact built-in types
+    # keep this comparison from calling user-defined equality methods.
+    return (all(type(value) is str or type(value) is int or value is None
+                for value in fields + original)
+            and fields == original)
+
+
 class TracebackException:
     """An exception ready for rendering.
 
@@ -1179,11 +1200,15 @@ def __init__(self, exc_type, exc_value, exc_traceback, *, limit=None,
             self.msg = exc_value.msg
             self._is_syntax_error = True
             metadata = getattr(exc_value, "_metadata", None)
-            if isinstance(metadata, tuple) and len(metadata) == 6:
-                (suggestion, self.text, lineno, self.offset,
-                 end_lineno, self.end_offset) = metadata
-                self.lineno = str(lineno)
-                self.end_lineno = str(end_lineno)
+            fields = (self.filename, self.msg, self.text, lno, self.offset,
+                      end_lno, self.end_offset)
+            if _syntax_error_metadata_matches(metadata, fields):
+                suggestion, location, original = metadata
+                if location is not None:
+                    (self.text, lineno, self.offset,
+                     end_lineno, self.end_offset) = location
+                    self.lineno = str(lineno)
+                    self.end_lineno = str(end_lineno)
                 self.msg = f"invalid syntax. Did you mean '{suggestion}'?"
         elif exc_type and issubclass(exc_type, ImportError) and \
                 getattr(exc_value, "name_from", None) is not None:

Parser/pegen.c

@@ -1108,9 +1108,6 @@ add_keyword_typo(KeywordTypo *candidates, int *count, int mark,
     if (word == NULL) {
         return -1;
     }
-    if (PyUnicode_GET_LENGTH(name) <= 2) {
-        max_cost = Py_MIN(max_cost, 2);
-    }
     Py_ssize_t cost = _Py_UTF8_Edit_Cost(name, word, max_cost);
     // Adjacent transpositions are one spelling error. The shared Levenshtein
     // distance counts them as two substitutions.
@@ -1181,17 +1178,20 @@ keyword_trial_accepts(Parser *p, const KeywordTypo *candidate)
 static int
 try_keyword_typo(Parser *p, const KeywordTypo *candidate, KeywordTrial *state)
 {
+    PyArena *arena = _PyArena_New();
+    if (arena == NULL) {
+        return -1;
+    }
     // Token text belongs to the original arena. Token headers and memo entries
     // must be private: even a failed trial can memoize nodes containing its edit.
-    Token *tokens = PyMem_Malloc(p->fill * sizeof(*tokens));
-    Token **index = PyMem_Malloc(p->fill * sizeof(*index));
+    Token *tokens = _PyArena_Malloc(arena, p->fill * sizeof(*tokens));
+    Token **index = _PyArena_Malloc(arena, p->fill * sizeof(*index));
     Parser trial = *p;
     trial.identifier_cache = p->identifier_cache == NULL ? NULL :
         PyMem_Malloc(IDENTIFIER_CACHE_SIZE * sizeof(*trial.identifier_cache));
-    PyArena *arena = _PyArena_New();
     PyObject *replacement = PyBytes_FromString(candidate->keyword);
     int accepted = -1;
-    if (tokens == NULL || index == NULL || arena == NULL || replacement == NULL ||
+    if (tokens == NULL || index == NULL || replacement == NULL ||
         (p->identifier_cache != NULL && trial.identifier_cache == NULL))
     {
         goto done;
@@ -1235,14 +1235,10 @@ try_keyword_typo(Parser *p, const KeywordTypo *candidate, KeywordTrial *state)
     accepted = keyword_trial_accepts(&trial, candidate);
 done:
     // No Parser_Free: the trial borrows the tokenizer and type-ignore comments,
-    // and its token headers are a single allocation.
+    // and its token arrays belong to the arena.
     PyMem_Free(trial.identifier_cache);
-    if (arena != NULL) {
-        _PyArena_Free(arena);
-    }
+    _PyArena_Free(arena);
     Py_XDECREF(replacement);
-    PyMem_Free(index);
-    PyMem_Free(tokens);
     return PyErr_Occurred() ? -1 : accepted;
 }
 
@@ -1356,7 +1352,8 @@ collect_keyword_typos(Parser *p, PySyntaxErrorObject *error, KeywordTypo *candid
         // Preserve missing-comma diagnostics and correctly spelled keywords
         // unless one edit repairs the input. For keywords, only substitute a
         // character: deleting one from "for" must not suggest "or".
-        Py_ssize_t max_cost = p->missing_comma || is_keyword ? 2 : 4;
+        Py_ssize_t max_cost =
+            p->missing_comma || is_keyword || PyUnicode_GET_LENGTH(name) <= 2 ? 2 : 4;
         int result = 0;
         for (int length = 0; length < p->n_keyword_lists && result == 0; length++) {
             if (is_keyword && length != text_length) {
@@ -1383,32 +1380,43 @@ collect_keyword_typos(Parser *p, PySyntaxErrorObject *error, KeywordTypo *candid
 }
 
 static PyObject *
-build_keyword_typo_metadata(Parser *p, const KeywordTypo *candidate)
+build_keyword_typo_metadata(Parser *p, PySyntaxErrorObject *error,
+                            const KeywordTypo *candidate)
 {
     Token *token = p->tokens[candidate->mark];
-    PyObject *line;
+    PyObject *line = NULL;
     if (_PyTokenizer_RetainedSource(p->tok) != NULL) {
         Py_ssize_t size;
         const char *text = _PyTokenizer_LineView(p->tok, token->lineno, &size);
         line = PyUnicode_DecodeUTF8(text, size, "replace");
     }
-    else {
-        _PyTokenizer_Info info = _PyTokenizer_GetInfo(p->tok);
-        line = _PyErr_ProgramDecodedTextObject(
-            info.filename, token->lineno, info.encoding);
+    else if (PyLong_AsLong(error->lineno) == token->lineno &&
+             error->text != NULL && PyUnicode_CheckExact(error->text)) {
+        // Reuse the error's line without another file access or codec lookup.
+        line = Py_NewRef(error->text);
     }
-    if (line == NULL) {
+    if (line == NULL && PyErr_Occurred()) {
         return NULL;
     }
-    Py_ssize_t start = _PyPegen_byte_offset_to_character_offset(line, token->col_offset);
-    Py_ssize_t end = _PyPegen_byte_offset_to_character_offset(line, token->end_col_offset);
-    PyObject *metadata = NULL;
-    if (start >= 0 && end >= 0) {
-        metadata = Py_BuildValue("(sOinin)", candidate->keyword, line,
-                                 token->lineno, start + 1,
-                                 token->end_lineno, end + 1);
+    PyObject *location = NULL;
+    if (line != NULL) {
+        Py_ssize_t start = _PyPegen_byte_offset_to_character_offset(line, token->col_offset);
+        Py_ssize_t end = _PyPegen_byte_offset_to_character_offset(line, token->end_col_offset);
+        if (start >= 0 && end >= 0) {
+            location = Py_BuildValue("(Oinin)", line, token->lineno, start + 1,
+                                    token->end_lineno, end + 1);
+        }
+        Py_DECREF(line);
+        if (location == NULL) {
+            return NULL;
+        }
     }
-    Py_DECREF(line);
+    // Without the typo's source line, retain the original error location.
+    PyObject *metadata = Py_BuildValue(
+        "(sO(OOOOOOO))", candidate->keyword, location != NULL ? location : Py_None,
+        error->filename, error->msg, error->text, error->lineno, error->offset,
+        error->end_lineno, error->end_offset);
+    Py_XDECREF(location);
     return metadata;
 }
 
@@ -1449,7 +1457,7 @@ keyword_typo_metadata(Parser *p, PySyntaxErrorObject *error)
             return NULL;
         }
         if (accepted) {
-            return build_keyword_typo_metadata(p, &candidates[i]);
+            return build_keyword_typo_metadata(p, error, &candidates[i]);
         }
     }
     return NULL;

Parser/pegen_errors.c

@@ -175,7 +175,7 @@ _PyPegen_tokenize_full_source_to_check_for_errors(Parser *p) {
         // Retain the tokens this error check already reads for bounded typo
         // trials. Never treat a truncated cache as complete input.
         if (cache_tokens && p->fill < MAX_KEYWORD_TYPO_TOKENS &&
-            (p->fill == 0 || p->tokens[p->fill - 1]->type != ENDMARKER) &&
+            p->tokens[p->fill - 1]->type != ENDMARKER &&
             new_token.type != ERRORTOKEN && new_token.type != TYPE_IGNORE &&
             _PyPegen_append_token(p, &new_token) < 0)
         {
@@ -203,7 +203,7 @@ _PyPegen_tokenize_full_source_to_check_for_errors(Parser *p) {
             }
             case ENDMARKER:
                 if (cache_tokens && p->fill < MAX_KEYWORD_TYPO_TOKENS &&
-                    p->start_rule == Py_single_input && p->fill > 0 &&
+                    p->start_rule == Py_single_input &&
                     p->tokens[p->fill - 1]->type == NEWLINE)
                 {
                     // append_token inserted single input's implicit newline;