bpo-43950: implement on-the-fly source tracking for interactive mode

Include/internal/pycore_parser.h

@@ -24,6 +24,17 @@ extern struct _mod* _PyParser_ASTFromFile(
     PyCompilerFlags *flags,
     int *errcode,
     PyArena *arena);
+extern struct _mod* _PyParser_InteractiveASTFromFile(
+    FILE *fp,
+    PyObject *filename_ob,
+    const char *enc,
+    int mode,
+    const char *ps1,
+    const char *ps2,
+    PyCompilerFlags *flags,
+    int *errcode,
+    PyObject **interactive_src,
+    PyArena *arena);
 
 #ifdef __cplusplus
 }

Lib/test/test_repl.py

@@ -107,6 +107,72 @@ def test_close_stdin(self):
         self.assertEqual(process.returncode, 0)
         self.assertIn('before close', output)
 
+    def test_interactive_source_tracking(self):
+        import ast
+
+        p = spawn_repl()
+        p.stdin.write(dedent("""\
+            import sys
+
+            sources = []
+            def source_tracker(frame, event, arg):
+                global sources
+                if source := frame.f_globals.get("__source__"):
+                    sources.append(source)
+
+            sys.settrace(source_tracker)
+
+            # basic statement
+            a = 1
+
+            # multiline statement
+            if a:
+                pass
+            else:
+                ...
+
+            # multiline expression
+            maybe = [
+                1,
+                2,
+                [
+                    3,
+                    4
+                ]
+            ][2][
+                0
+            ]
+
+            # basic expression
+            sys.settrace(None)
+
+            print("\\n" + repr(sources) + "\\n")
+        """))
+        output = kill_python(p)
+        result_line = output.splitlines()[-3]
+        tracing_records = ast.literal_eval(result_line)
+        self.assertIn("a = 1\n", tracing_records)
+        self.assertIn("sys.settrace(None)\n", tracing_records)
+        self.assertIn("if a:\n    pass\nelse:\n    ...\n\n", tracing_records)
+        self.assertIn("maybe = [\n    1,\n    2,\n    [\n        3,\n"
+                      "        4\n    ]\n][2][\n    0\n]\n", tracing_records)
+
+    def test_interactive_traceback_reporting(self):
+        user_input = "1 / 0 / 3 / 4"
+        p = spawn_repl()
+        p.stdin.write(user_input)
+        output = kill_python(p)
+        self.assertEqual(p.returncode, 0)
+
+        traceback_lines = output.splitlines()[-6:-1]
+        expected_lines = [
+            "Traceback (most recent call last):",
+            "  File \"<stdin>\", line 1, in <module>",
+            "    1 / 0 / 3 / 4",
+            "    ~~^~~",
+            "ZeroDivisionError: division by zero",
+        ]
+        self.assertEqual(traceback_lines, expected_lines)
 
 if __name__ == "__main__":
     unittest.main()

Parser/peg_api.c

@@ -24,5 +24,19 @@ _PyParser_ASTFromFile(FILE *fp, PyObject *filename_ob, const char *enc,
         return NULL;
     }
     return _PyPegen_run_parser_from_file_pointer(fp, mode, filename_ob, enc, ps1, ps2,
-                                        flags, errcode, arena);
+                                                 flags, errcode, NULL, arena);
+}
+
+
+mod_ty
+_PyParser_InteractiveASTFromFile(FILE *fp, PyObject *filename_ob, const char *enc,
+                                 int mode, const char *ps1, const char* ps2,
+                                 PyCompilerFlags *flags, int *errcode,
+                                 PyObject **interactive_src, PyArena *arena)
+{
+    if (PySys_Audit("compile", "OO", Py_None, filename_ob) < 0) {
+        return NULL;
+    }
+    return _PyPegen_run_parser_from_file_pointer(fp, mode, filename_ob, enc, ps1, ps2,
+                                                 flags, errcode, interactive_src, arena);
 }

Parser/pegen.c

@@ -1373,7 +1373,8 @@ _PyPegen_run_parser(Parser *p)
 mod_ty
 _PyPegen_run_parser_from_file_pointer(FILE *fp, int start_rule, PyObject *filename_ob,
                              const char *enc, const char *ps1, const char *ps2,
-                             PyCompilerFlags *flags, int *errcode, PyArena *arena)
+                             PyCompilerFlags *flags, int *errcode,
+                             PyObject **interactive_src, PyArena *arena)
 {
     struct tok_state *tok = PyTokenizer_FromFile(fp, enc, ps1, ps2);
     if (tok == NULL) {
@@ -1404,6 +1405,15 @@ _PyPegen_run_parser_from_file_pointer(FILE *fp, int start_rule, PyObject *filena
     result = _PyPegen_run_parser(p);
     _PyPegen_Parser_Free(p);
 
+    if (tok->fp_interactive && tok->interactive_src_start && result && interactive_src != NULL) {
+        *interactive_src = PyUnicode_FromString(tok->interactive_src_start);
+        if (!interactive_src || _PyArena_AddPyObject(arena, *interactive_src) < 0) {
+            Py_XDECREF(interactive_src);
+            result = NULL;
+            goto error;
+        }
+    }
+
 error:
     PyTokenizer_Free(tok);
     return result;

Parser/pegen.h

@@ -250,7 +250,8 @@ PyObject *_PyPegen_new_identifier(Parser *, const char *);
 Parser *_PyPegen_Parser_New(struct tok_state *, int, int, int, int *, PyArena *);
 void _PyPegen_Parser_Free(Parser *);
 mod_ty _PyPegen_run_parser_from_file_pointer(FILE *, int, PyObject *, const char *,
-                                    const char *, const char *, PyCompilerFlags *, int *, PyArena *);
+                                    const char *, const char *, PyCompilerFlags *, int *, PyObject **,
+                                    PyArena *);
 void *_PyPegen_run_parser(Parser *);
 mod_ty _PyPegen_run_parser_from_string(const char *, int, PyObject *, PyCompilerFlags *, PyArena *);
 asdl_stmt_seq *_PyPegen_interactive_exit(Parser *);

Python/pythonrun.c

@@ -253,8 +253,9 @@ PyRun_InteractiveOneObjectEx(FILE *fp, PyObject *filename,
         return -1;
     }
 
-    mod = _PyParser_ASTFromFile(fp, filename, enc, Py_single_input,
-                                ps1, ps2, flags, &errcode, arena);
+    PyObject *interactive_src = NULL;
+    mod = _PyParser_InteractiveASTFromFile(fp, filename, enc, Py_single_input,
+                                           ps1, ps2, flags, &errcode, &interactive_src, arena);
 
     Py_XDECREF(v);
     Py_XDECREF(w);
@@ -273,6 +274,12 @@ PyRun_InteractiveOneObjectEx(FILE *fp, PyObject *filename,
         return -1;
     }
     d = PyModule_GetDict(m);
+    if (interactive_src) {
+        if (PyDict_SetItemString(d, "__source__", interactive_src) < 0) {
+            _PyArena_Free(arena);
+            return -1;
+        }
+    }
     v = run_mod(mod, filename, d, d, flags, arena);
     _PyArena_Free(arena);
     if (v == NULL) {

Python/traceback.c

@@ -375,6 +375,62 @@ _Py_FindSourceFile(PyObject *filename, char* namebuf, size_t namelen, PyObject *
     return result;
 }
 
+int
+_Py_DisplayLine(PyObject *f, PyObject *line_obj, PyObject **line,
+                int indent, int *truncation)
+{
+    if (line) {
+        Py_INCREF(line_obj);
+        *line = line_obj;
+    }
+
+    int i, kind, err = -1;
+    const void *data;
+    char buf[MAXPATHLEN+1];
+
+    /* remove the indentation of the line */
+    kind = PyUnicode_KIND(line_obj);
+    data = PyUnicode_DATA(line_obj);
+    for (i=0; i < PyUnicode_GET_LENGTH(line_obj); i++) {
+        Py_UCS4 ch = PyUnicode_READ(kind, data, i);
+        if (ch != ' ' && ch != '\t' && ch != '\014')
+            break;
+    }
+    if (i) {
+        PyObject *truncated;
+        truncated = PyUnicode_Substring(line_obj, i, PyUnicode_GET_LENGTH(line_obj));
+        if (truncated) {
+            Py_DECREF(line_obj);
+            line_obj = truncated;
+        } else {
+            PyErr_Clear();
+        }
+    }
+
+    if (truncation != NULL) {
+        *truncation = i - indent;
+    }
+
+    /* Write some spaces before the line */
+    strcpy(buf, "          ");
+    assert (strlen(buf) == 10);
+    while (indent > 0) {
+        if (indent < 10)
+            buf[indent] = '\0';
+        err = PyFile_WriteString(buf, f);
+        if (err != 0)
+            break;
+        indent -= 10;
+    }
+
+    /* finally display the line */
+    if (err == 0)
+        err = PyFile_WriteObject(line_obj, f, Py_PRINT_RAW);
+    Py_DECREF(line_obj);
+    if  (err == 0)
+        err = PyFile_WriteString("\n", f);
+    return err;
+}
 int
 _Py_DisplaySourceLine(PyObject *f, PyObject *filename, int lineno, int indent, int *truncation, PyObject **line)
 {
@@ -389,8 +445,6 @@ _Py_DisplaySourceLine(PyObject *f, PyObject *filename, int lineno, int indent, i
     PyObject *lineobj = NULL;
     PyObject *res;
     char buf[MAXPATHLEN+1];
-    int kind;
-    const void *data;
 
     /* open the file */
     if (filename == NULL)
@@ -467,53 +521,34 @@ _Py_DisplaySourceLine(PyObject *f, PyObject *filename, int lineno, int indent, i
         return err;
     }
 
-    if (line) {
-        Py_INCREF(lineobj);
-        *line = lineobj;
-    }
+    return _Py_DisplayLine(f, lineobj, line, indent, truncation);
+}
 
-    /* remove the indentation of the line */
-    kind = PyUnicode_KIND(lineobj);
-    data = PyUnicode_DATA(lineobj);
-    for (i=0; i < PyUnicode_GET_LENGTH(lineobj); i++) {
-        Py_UCS4 ch = PyUnicode_READ(kind, data, i);
-        if (ch != ' ' && ch != '\t' && ch != '\014')
-            break;
-    }
-    if (i) {
-        PyObject *truncated;
-        truncated = PyUnicode_Substring(lineobj, i, PyUnicode_GET_LENGTH(lineobj));
-        if (truncated) {
-            Py_DECREF(lineobj);
-            lineobj = truncated;
-        } else {
-            PyErr_Clear();
-        }
+int
+_Py_DisplayInteractiveSourceLine(PyObject *f, PyFrameObject *frame, int lineno, int indent,
+                                 int *truncation, PyObject **line)
+{
+    PyObject *globals = _PyFrame_GetGlobals(frame);
+    PyObject *source = PyDict_GetItemString(globals, "__source__");
+    if (!source) {
+        return -1;
     }
 
-    if (truncation != NULL) {
-        *truncation = i - indent;
+    PyObject *lines = PyUnicode_Splitlines(source, 0);
+    if (!lines || PyList_GET_SIZE(lines) < lineno) {
+        Py_XDECREF(lines);
+        return -1;
     }
 
-    /* Write some spaces before the line */
-    strcpy(buf, "          ");
-    assert (strlen(buf) == 10);
-    while (indent > 0) {
-        if (indent < 10)
-            buf[indent] = '\0';
-        err = PyFile_WriteString(buf, f);
-        if (err != 0)
-            break;
-        indent -= 10;
+    PyObject *lineobj = PyList_GetItem(lines, lineno - 1);
+    if (!lineobj) {
+        Py_DECREF(lines);
+        return -1;
     }
+    Py_INCREF(lineobj);
+    Py_DECREF(lines);
 
-    /* finally display the line */
-    if (err == 0)
-        err = PyFile_WriteObject(lineobj, f, Py_PRINT_RAW);
-    Py_DECREF(lineobj);
-    if  (err == 0)
-        err = PyFile_WriteString("\n", f);
-    return err;
+    return _Py_DisplayLine(f, lineobj, line, indent, truncation);
 }
 
 /* AST based Traceback Specialization
@@ -702,8 +737,16 @@ tb_displayline(PyTracebackObject* tb, PyObject *f, PyObject *filename, int linen
     int truncation = _TRACEBACK_SOURCE_LINE_INDENT;
     PyObject* source_line = NULL;
 
-    if (_Py_DisplaySourceLine(f, filename, lineno, _TRACEBACK_SOURCE_LINE_INDENT,
-                               &truncation, &source_line) != 0) {
+    if (PyUnicode_CompareWithASCIIString(filename, "<stdin>") == 0) {
+        err = _Py_DisplayInteractiveSourceLine(f, frame, lineno, _TRACEBACK_SOURCE_LINE_INDENT,
+                                               &truncation, &source_line);
+    }
+    else {
+        err = _Py_DisplaySourceLine(f, filename, lineno, _TRACEBACK_SOURCE_LINE_INDENT,
+                                    &truncation, &source_line);
+    }
+
+    if (err != 0) {
         /* ignore errors since we can't report them, can we? */
         err = ignore_source_errors();
         goto done;

Tools/peg_generator/peg_extension/peg_extension.c

@@ -53,7 +53,7 @@ parse_file(PyObject *self, PyObject *args, PyObject *kwds)
     PyCompilerFlags flags = _PyCompilerFlags_INIT;
     mod_ty res = _PyPegen_run_parser_from_file_pointer(
                         fp, Py_file_input, filename_ob,
-                        NULL, NULL, NULL, &flags, NULL, arena);
+                        NULL, NULL, NULL, &flags, NULL, NULL, arena);
     fclose(fp);
     if (res == NULL) {
         goto error;