bpo-43950: Add code.co_positions (PEP 657) (GH-26955)

This PR is part of PEP 657 and augments the compiler to emit ending
line numbers as well as starting and ending columns from the AST
into compiled code objects. This allows bytecodes to be correlated
to the exact source code ranges that generated them.

This information is made available through the following public APIs:

* The `co_positions` method on code objects.
* The C API function `PyCode_Addr2Location`.

Co-authored-by: Batuhan Taskaya <isidentical@gmail.com>
Co-authored-by: Ammar Askar <ammar@ammaraskar.com>

1 files changed, 83 insertions, 0 deletions

diff --git a/Lib/test/test_code.py b/Lib/test/test_code.py
index 27342ad..55ba30a 100644
--- a/Lib/test/test_code.py
+++ b/Lib/test/test_code.py
@@ -129,7 +129,9 @@ import inspect
 import sys
 import threading
 import unittest
+import textwrap
 import weakref
+
 try:
     import ctypes
 except ImportError:
@@ -223,6 +225,8 @@ class CodeTest(unittest.TestCase):
                         co.co_name,
                         co.co_firstlineno,
                         co.co_lnotab,
+                        co.co_endlinetable,
+                        co.co_columntable,
                         co.co_exceptiontable,
                         co.co_freevars,
                         co.co_cellvars)
@@ -257,6 +261,8 @@ class CodeTest(unittest.TestCase):
             ("co_filename", "newfilename"),
             ("co_name", "newname"),
             ("co_linetable", code2.co_linetable),
+            ("co_endlinetable", code2.co_endlinetable),
+            ("co_columntable", code2.co_columntable),
         ):
             with self.subTest(attr=attr, value=value):
                 new_code = code.replace(**{attr: value})
@@ -293,6 +299,8 @@ class CodeTest(unittest.TestCase):
                          co.co_name,
                          co.co_firstlineno,
                          co.co_lnotab,
+                         co.co_endlinetable,
+                         co.co_columntable,
                          co.co_exceptiontable,
                          co.co_freevars,
                          co.co_cellvars,
@@ -309,6 +317,81 @@ class CodeTest(unittest.TestCase):
         new_code = code = func.__code__.replace(co_linetable=b'')
         self.assertEqual(list(new_code.co_lines()), [])
 
+    def test_co_positions_artificial_instructions(self):
+        import dis
+
+        namespace = {}
+        exec(textwrap.dedent("""\
+        try:
+            1/0
+        except Exception as e:
+            exc = e
+        """), namespace)
+
+        exc = namespace['exc']
+        traceback = exc.__traceback__
+        code = traceback.tb_frame.f_code
+
+        artificial_instructions = []
+        for instr, positions in zip(
+            dis.get_instructions(code),
+            code.co_positions(),
+            strict=True
+        ):
+            # If any of the positions is None, then all have to
+            # be None as well for the case above. There are still
+            # some places in the compiler, where the artificial instructions
+            # get assigned the first_lineno but they don't have other positions.
+            # There is no easy way of inferring them at that stage, so for now
+            # we don't support it.
+            self.assertTrue(all(positions) or not any(positions))
+
+            if not any(positions):
+                artificial_instructions.append(instr)
+
+        self.assertEqual(
+            [
+                (instruction.opname, instruction.argval)
+                for instruction in artificial_instructions
+            ],
+            [
+                ("PUSH_EXC_INFO", None),
+                ("LOAD_CONST", None), # artificial 'None'
+                ("STORE_NAME", "e"),  # XX: we know the location for this
+                ("DELETE_NAME", "e"),
+                ("RERAISE", 1),
+                ("POP_EXCEPT_AND_RERAISE", None)
+            ]
+        )
+
+    # co_positions behavior when info is missing.
+
+    def test_co_positions_empty_linetable(self):
+        def func():
+            x = 1
+        new_code = func.__code__.replace(co_linetable=b'')
+        for line, end_line, column, end_column in new_code.co_positions():
+            self.assertIsNone(line)
+            self.assertEqual(end_line, new_code.co_firstlineno + 1)
+
+    def test_co_positions_empty_endlinetable(self):
+        def func():
+            x = 1
+        new_code = func.__code__.replace(co_endlinetable=b'')
+        for line, end_line, column, end_column in new_code.co_positions():
+            self.assertEqual(line, new_code.co_firstlineno + 1)
+            self.assertIsNone(end_line)
+
+    def test_co_positions_empty_columntable(self):
+        def func():
+            x = 1
+        new_code = func.__code__.replace(co_columntable=b'')
+        for line, end_line, column, end_column in new_code.co_positions():
+            self.assertEqual(line, new_code.co_firstlineno + 1)
+            self.assertEqual(end_line, new_code.co_firstlineno + 1)
+            self.assertIsNone(column)
+            self.assertIsNone(end_column)
+
 
 def isinterned(s):
     return s is sys.intern(('_' + s + '_')[1:-1])