gh-110805: Allow the repl to show source code and complete tracebacks (#110775)

11 files changed, 191 insertions, 19 deletions

diff --git a/Include/internal/pycore_interp.h b/Include/internal/pycore_interp.h
index 80f16af..60d333a 100644
--- a/Include/internal/pycore_interp.h
+++ b/Include/internal/pycore_interp.h
@@ -233,6 +233,7 @@ struct _is {
 
    /* the initial PyInterpreterState.threads.head */
     PyThreadState _initial_thread;
+    Py_ssize_t _interactive_src_count;
 };
 
 
diff --git a/Include/internal/pycore_parser.h b/Include/internal/pycore_parser.h
index dd51b92..067b34c 100644
--- a/Include/internal/pycore_parser.h
+++ b/Include/internal/pycore_parser.h
@@ -58,7 +58,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
 }
diff --git a/Lib/linecache.py b/Lib/linecache.py
index 97644a8..c1c988d 100644
--- a/Lib/linecache.py
+++ b/Lib/linecache.py
@@ -180,3 +180,10 @@ def lazycache(filename, module_globals):
             cache[filename] = (get_lines,)
             return True
     return False
+
+def _register_code(code, string, name):
+    cache[code] = (
+            len(string),
+            None,
+            [line + '\n' for line in string.splitlines()],
+            name)
diff --git a/Lib/test/test_cmd_line_script.py b/Lib/test/test_cmd_line_script.py
index 1b58882..614c6b3 100644
--- a/Lib/test/test_cmd_line_script.py
+++ b/Lib/test/test_cmd_line_script.py
@@ -203,6 +203,8 @@ class CmdLineTest(unittest.TestCase):
             stderr = p.stderr if separate_stderr else p.stdout
             self.assertIn(b'Traceback ', stderr.readline())
             self.assertIn(b'File "<stdin>"', stderr.readline())
+            self.assertIn(b'1/0', stderr.readline())
+            self.assertIn(b'    ~^~', stderr.readline())
             self.assertIn(b'ZeroDivisionError', stderr.readline())
 
     def test_repl_stdout_flush(self):
diff --git a/Lib/test/test_repl.py b/Lib/test/test_repl.py
index 58392f2..2ee5117 100644
--- a/Lib/test/test_repl.py
+++ b/Lib/test/test_repl.py
@@ -131,6 +131,68 @@ class TestInteractiveInterpreter(unittest.TestCase):
         self.assertEqual(process.returncode, 0)
         self.assertIn('before close', output)
 
+    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)
+
+    def test_interactive_traceback_reporting_multiple_input(self):
+        user_input1 = dedent("""
+        def foo(x):
+            1 / x
+
+        """)
+        p = spawn_repl()
+        p.stdin.write(user_input1)
+        user_input2 = "foo(0)"
+        p.stdin.write(user_input2)
+        output = kill_python(p)
+        self.assertEqual(p.returncode, 0)
+
+        traceback_lines = output.splitlines()[-7:-1]
+        expected_lines = [
+            '  File "<stdin>", line 1, in <module>',
+            '    foo(0)',
+            '  File "<stdin>", line 2, in foo',
+            '    1 / x',
+            '    ~~^~~',
+            'ZeroDivisionError: division by zero'
+        ]
+        self.assertEqual(traceback_lines, expected_lines)
+
+    def test_interactive_source_is_in_linecache(self):
+        user_input = dedent("""
+        def foo(x):
+            return x + 1
+
+        def bar(x):
+            return foo(x) + 2
+        """)
+        p = spawn_repl()
+        p.stdin.write(user_input)
+        user_input2 = dedent("""
+        import linecache
+        print(linecache.cache['<python-input-1>'])
+        """)
+        p.stdin.write(user_input2)
+        output = kill_python(p)
+        self.assertEqual(p.returncode, 0)
+        expected = "(30, None, [\'def foo(x):\\n\', \'    return x + 1\\n\', \'\\n\'], \'<stdin>\')"
+        self.assertIn(expected, output, expected)
+
+
 
 class TestInteractiveModeSyntaxErrors(unittest.TestCase):
 
diff --git a/Lib/traceback.py b/Lib/traceback.py
index 12fcdad..7cc84b9 100644
--- a/Lib/traceback.py
+++ b/Lib/traceback.py
@@ -434,7 +434,6 @@ class StackSummary(list):
             co = f.f_code
             filename = co.co_filename
             name = co.co_name
-
             fnames.add(filename)
             linecache.lazycache(filename, f.f_globals)
             # Must defer line lookups until we have called checkcache.
@@ -447,6 +446,7 @@ class StackSummary(list):
                 end_lineno=end_lineno, colno=colno, end_colno=end_colno))
         for filename in fnames:
             linecache.checkcache(filename)
+
         # If immediate lookup was desired, trigger lookups now.
         if lookup_lines:
             for f in result:
@@ -479,8 +479,12 @@ class StackSummary(list):
         gets called for every frame to be printed in the stack summary.
         """
         row = []
-        row.append('  File "{}", line {}, in {}\n'.format(
-            frame_summary.filename, frame_summary.lineno, frame_summary.name))
+        if frame_summary.filename.startswith("<python-input"):
+            row.append('  File "<stdin>", line {}, in {}\n'.format(
+                frame_summary.lineno, frame_summary.name))
+        else:
+            row.append('  File "{}", line {}, in {}\n'.format(
+                frame_summary.filename, frame_summary.lineno, frame_summary.name))
         if frame_summary.line:
             stripped_line = frame_summary.line.strip()
             row.append('    {}\n'.format(stripped_line))
diff --git a/Misc/NEWS.d/next/Core and Builtins/2023-10-13-09-21-29.gh-issue-110805.vhU7A7.rst b/Misc/NEWS.d/next/Core and Builtins/2023-10-13-09-21-29.gh-issue-110805.vhU7A7.rst
new file mode 100644
index 0000000..be90bb3
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2023-10-13-09-21-29.gh-issue-110805.vhU7A7.rst
@@ -0,0 +1,2 @@
+Allow the repl to show source code and complete tracebacks. Patch by Pablo
+Galindo
diff --git a/Parser/peg_api.c b/Parser/peg_api.c
index 447f566..d4acc3e 100644
--- a/Parser/peg_api.c
+++ b/Parser/peg_api.c
@@ -23,5 +23,18 @@ _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);
+}
\ No newline at end of file
diff --git a/Parser/pegen.c b/Parser/pegen.c
index 2e0b569..0c60394 100644
--- a/Parser/pegen.c
+++ b/Parser/pegen.c
@@ -878,7 +878,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) {
@@ -908,6 +909,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;
diff --git a/Parser/pegen.h b/Parser/pegen.h
index 266d521..424f80a 100644
--- a/Parser/pegen.h
+++ b/Parser/pegen.h
@@ -350,7 +350,8 @@ void *_PyPegen_nonparen_genexp_in_call(Parser *p, expr_ty args, asdl_comprehensi
 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 *);
diff --git a/Python/pythonrun.c b/Python/pythonrun.c
index 2e27471..7499429 100644
--- a/Python/pythonrun.c
+++ b/Python/pythonrun.c
@@ -40,7 +40,7 @@
 /* Forward */
 static void flush_io(void);
 static PyObject *run_mod(mod_ty, PyObject *, PyObject *, PyObject *,
-                          PyCompilerFlags *, PyArena *);
+                          PyCompilerFlags *, PyArena *, PyObject*);
 static PyObject *run_pyc_file(FILE *, PyObject *, PyObject *,
                               PyCompilerFlags *);
 static int PyRun_InteractiveOneObjectEx(FILE *, PyObject *, PyCompilerFlags *);
@@ -178,7 +178,8 @@ PyRun_InteractiveLoopFlags(FILE *fp, const char *filename, PyCompilerFlags *flag
 // Call _PyParser_ASTFromFile() with sys.stdin.encoding, sys.ps1 and sys.ps2
 static int
 pyrun_one_parse_ast(FILE *fp, PyObject *filename,
-                    PyCompilerFlags *flags, PyArena *arena, mod_ty *pmod)
+                    PyCompilerFlags *flags, PyArena *arena,
+                    mod_ty *pmod, PyObject** interactive_src)
 {
     PyThreadState *tstate = _PyThreadState_GET();
 
@@ -236,9 +237,9 @@ pyrun_one_parse_ast(FILE *fp, PyObject *filename,
     }
 
     int errcode = 0;
-    *pmod = _PyParser_ASTFromFile(fp, filename, encoding,
-                                  Py_single_input, ps1, ps2,
-                                  flags, &errcode, arena);
+    *pmod = _PyParser_InteractiveASTFromFile(fp, filename, encoding,
+                                             Py_single_input, ps1, ps2,
+                                             flags, &errcode, interactive_src, arena);
     Py_XDECREF(ps1_obj);
     Py_XDECREF(ps2_obj);
     Py_XDECREF(encoding_obj);
@@ -266,7 +267,8 @@ PyRun_InteractiveOneObjectEx(FILE *fp, PyObject *filename,
     }
 
     mod_ty mod;
-    int parse_res = pyrun_one_parse_ast(fp, filename, flags, arena, &mod);
+    PyObject *interactive_src;
+    int parse_res = pyrun_one_parse_ast(fp, filename, flags, arena, &mod, &interactive_src);
     if (parse_res != 0) {
         _PyArena_Free(arena);
         return parse_res;
@@ -279,7 +281,7 @@ PyRun_InteractiveOneObjectEx(FILE *fp, PyObject *filename,
     }
     PyObject *main_dict = PyModule_GetDict(main_module);  // borrowed ref
 
-    PyObject *res = run_mod(mod, filename, main_dict, main_dict, flags, arena);
+    PyObject *res = run_mod(mod, filename, main_dict, main_dict, flags, arena, interactive_src);
     _PyArena_Free(arena);
     Py_DECREF(main_module);
     if (res == NULL) {
@@ -1149,7 +1151,7 @@ PyRun_StringFlags(const char *str, int start, PyObject *globals,
             str, &_Py_STR(anon_string), start, flags, arena);
 
     if (mod != NULL)
-        ret = run_mod(mod, &_Py_STR(anon_string), globals, locals, flags, arena);
+        ret = run_mod(mod, &_Py_STR(anon_string), globals, locals, flags, arena, NULL);
     _PyArena_Free(arena);
     return ret;
 }
@@ -1174,7 +1176,7 @@ pyrun_file(FILE *fp, PyObject *filename, int start, PyObject *globals,
 
     PyObject *ret;
     if (mod != NULL) {
-        ret = run_mod(mod, filename, globals, locals, flags, arena);
+        ret = run_mod(mod, filename, globals, locals, flags, arena, NULL);
     }
     else {
         ret = NULL;
@@ -1262,12 +1264,70 @@ run_eval_code_obj(PyThreadState *tstate, PyCodeObject *co, PyObject *globals, Py
 
 static PyObject *
 run_mod(mod_ty mod, PyObject *filename, PyObject *globals, PyObject *locals,
-            PyCompilerFlags *flags, PyArena *arena)
+            PyCompilerFlags *flags, PyArena *arena, PyObject* interactive_src)
 {
     PyThreadState *tstate = _PyThreadState_GET();
-    PyCodeObject *co = _PyAST_Compile(mod, filename, flags, -1, arena);
-    if (co == NULL)
+    PyObject* interactive_filename = filename;
+    if (interactive_src) {
+        PyInterpreterState *interp = tstate->interp;
+        interactive_filename = PyUnicode_FromFormat(
+            "<python-input-%d>", interp->_interactive_src_count++
+        );
+        if (interactive_filename == NULL) {
+            return NULL;
+        }
+    }
+
+    PyCodeObject *co = _PyAST_Compile(mod, interactive_filename, flags, -1, arena);
+    if (co == NULL) {
+        Py_DECREF(interactive_filename);
         return NULL;
+    }
+
+    if (interactive_src) {
+        PyObject *linecache_module = PyImport_ImportModule("linecache");
+
+        if (linecache_module == NULL) {
+            Py_DECREF(co);
+            Py_DECREF(interactive_filename);
+            return NULL;
+        }
+
+        PyObject *print_tb_func = PyObject_GetAttrString(linecache_module, "_register_code");
+
+        if (print_tb_func == NULL) {
+            Py_DECREF(co);
+            Py_DECREF(interactive_filename);
+            Py_DECREF(linecache_module);
+            return NULL;
+        }
+
+        if (!PyCallable_Check(print_tb_func)) {
+            Py_DECREF(co);
+            Py_DECREF(interactive_filename);
+            Py_DECREF(linecache_module);
+            Py_DECREF(print_tb_func);
+            PyErr_SetString(PyExc_ValueError, "linecache._register_code is not callable");
+            return NULL;
+        }
+
+        PyObject* result = PyObject_CallFunction(
+            print_tb_func, "OOO",
+            interactive_filename,
+            interactive_src,
+            filename
+        );
+
+        Py_DECREF(interactive_filename);
+
+        Py_DECREF(linecache_module);
+        Py_XDECREF(print_tb_func);
+        Py_XDECREF(result);
+        if (!result) {
+            Py_DECREF(co);
+            return NULL;
+        }
+    }
 
     if (_PySys_Audit(tstate, "exec", "O", co) < 0) {
         Py_DECREF(co);