gh-159672: Preserve error state during keyword replay

Author: pablogsal

Lib/test/test_traceback.py

@@ -1912,6 +1912,74 @@ 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):
+                with self.assertRaises(SyntaxError) as caught:
+                    compile("impor math", "<test>", "exec")
+                error = caught.exception
+                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")
+
+        with self.assertRaises(SyntaxError) as caught:
+            compile("impor math", "<test>", "exec")
+        error = caught.exception
+        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

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

@@ -1113,9 +1113,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.
@@ -1186,17 +1183,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;
@@ -1240,14 +1240,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;
 }
 
@@ -1358,7 +1354,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) {
@@ -1385,32 +1382,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;
 }
 
@@ -1451,7 +1459,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;