Add test for co_branches

python · markshannon · Jan 6, 2025 · Dec 31, 2024 · Dec 31, 2024 · Jan 3, 2025
commit 182c512200d7dd37af28e8c85d3e645f78e19d6b
diff --git a/Lib/test/test_code.py b/Lib/test/test_code.py
@@ -215,6 +215,8 @@
 from test.support import threading_helper, import_helper
 from test.support.bytecode_helper import instructions_with_positions
 from opcode import opmap, opname
+from _testcapi import code_offset_to_line
+
 COPY_FREE_VARS = opmap['COPY_FREE_VARS']
 
 
@@ -896,6 +898,44 @@ async def async_func():
 
         rc, out, err = assert_python_ok('-OO', '-c', code)
 
+    def test_co_branches(self):
+
+        def get_line_branches(func):
+            code = func.__code__
+            base = code.co_firstlineno
+            return [
+                (
+                    code_offset_to_line(code, src)-base,
+                    code_offset_to_line(code, left)-base,
+                    code_offset_to_line(code, right)-base
-                    code_offset_to_line(code, src)-base,
-                    code_offset_to_line(code, left)-base,
-                    code_offset_to_line(code, right)-base
+                    code_offset_to_line(code, src) - base,
+                    code_offset_to_line(code, left) - base,
+                    code_offset_to_line(code, right) - base,
-                    code_offset_to_line(code, src)-base,
-                    code_offset_to_line(code, left)-base,
-                    code_offset_to_line(code, right)-base
+                    code_offset_to_line(code, src) - base,
+                    code_offset_to_line(code, left) - base,
+                    code_offset_to_line(code, right) - base,
+                ) for (src, left, right) in
+                code.co_branches()
+            ]
+
+        def simple(x):
+            if x:
+                A
+            else:
+                B
+
+        self.assertEqual(
+            get_line_branches(simple),
+            [(1,2,4)])
+
+        def with_extended_args(x):
+            if x:
+                A.x; A.x; A.x; A.x; A.x; A.x;
+                A.x; A.x; A.x; A.x; A.x; A.x;
+                A.x; A.x; A.x; A.x; A.x; A.x;
+                A.x; A.x; A.x; A.x; A.x; A.x;
+                A.x; A.x; A.x; A.x; A.x; A.x;
+            else:
+                B
+
+        self.assertEqual(
+            get_line_branches(with_extended_args),
+            [(1,2,8)])
+
 if check_impl_detail(cpython=True) and ctypes is not None:
     py = ctypes.pythonapi
     freefunc = ctypes.CFUNCTYPE(None,ctypes.c_voidp)

diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
@@ -3415,6 +3415,26 @@ test_atexit(PyObject *self, PyObject *Py_UNUSED(args))
     Py_RETURN_NONE;
 }
 
+static PyObject*
+code_offset_to_line(PyObject* self, PyObject* const* args, Py_ssize_t nargsf)
+{
+    Py_ssize_t nargs = _PyVectorcall_NARGS(nargsf);
+    if (nargs != 2) {
+        PyErr_SetString(PyExc_TypeError, "code_offset_to_line takes 2 arguments");
+        return NULL;
+    }
+    int offset;
+    if (PyLong_AsInt32(args[1], &offset) < 0) {
+        return NULL;
+    }
+    PyCodeObject *code = (PyCodeObject *)args[0];
+    if (!PyCode_Check(code)) {
+        PyErr_SetString(PyExc_TypeError, "first arg must be a code object");
+        return NULL;
+    }
+    return PyLong_FromInt32(PyCode_Addr2Line(code, offset));
+}
+
 static PyMethodDef TestMethods[] = {
     {"set_errno",               set_errno,                       METH_VARARGS},
     {"test_config",             test_config,                     METH_NOARGS},
@@ -3557,6 +3577,7 @@ static PyMethodDef TestMethods[] = {
     {"finalize_thread_hang", finalize_thread_hang, METH_O, NULL},
     {"type_freeze", type_freeze, METH_VARARGS},
     {"test_atexit", test_atexit, METH_NOARGS},
+    {"code_offset_to_line", _PyCFunction_CAST(code_offset_to_line), METH_FASTCALL},
     {NULL, NULL} /* sentinel */
 };