cython: revamp towards 3.11

pulkin · Feb 3, 2024 · c816c88 · c816c88
1 parent b667415
commit c816c88
Show file tree

Hide file tree

Showing 3 changed files with 120 additions and 57 deletions.
diff --git a/cython/frame.pyx b/cython/frame.pyx
@@ -5,59 +5,115 @@ from .primitives import NULL as NULL_object, block_stack_item
 
 
 cdef extern from "frameobject.h":
-    ctypedef struct PyTryBlock:
-        int b_type
-        int b_handler
-        int b_level
-
-    struct _frame:
+    ctypedef struct PyFrameObject:
         PyObject** f_valuestack
-        PyTryBlock* f_blockstack
-        int f_iblock
         PyObject** f_localsplus
 
 
-cdef extern from *:  # stack depth for different python versions
+cdef extern from *:
     """
-    #if PY_VERSION_HEX >= 0x030A0000
-      static int _pyteleport_stackdepth(struct _frame* frame) {
-        return frame->f_stackdepth;
-      }
-    #elif PY_VERSION_HEX >= 0x03080000
-      static int _pyteleport_stackdepth(struct _frame* frame) {
-        if (frame->f_stacktop)
-          return (int) (frame->f_stacktop - frame->f_valuestack);
-        else
-          return -1;
-      }
-    #elif defined(PY_VERSION_HEX)
-      #error "Unknown python version"
+    #ifndef PY_VERSION_HEX
+        #error "PY_VERSION_HEX not defined"
+    #else
+        #define PYTELEPORT_PYTHON_VERSION (PY_VERSION_HEX >> 16)
+    #endif
+
+    #if PYTELEPORT_PYTHON_VERSION == 0x030B
+        #include "internal/pycore_frame.h"
+
+        #define FRAME (frame->f_frame)
+        #define CODE (FRAME->f_code)
+        static PyObject** _pyframe_get_value_stack(PyFrameObject* frame) {return FRAME->localsplus + CODE->co_nlocalsplus;}
+        static int _pyframe_get_value_stack_depth(PyFrameObject* frame) {return FRAME->stacktop;}
+
+        #define _PYFRAME_DEFINE_BLOCK_STACK_GETTER(name) static int _pyframe_get_block_stack_ ## name(PyFrameObject* frame, int i) {return -1;}
+        static int _pyframe_get_block_stack_depth(PyFrameObject* frame) {return -1;}
+
+        static int _pyframe_n_locals(PyFrameObject* frame) {return CODE->co_nlocals;}
+        static int _pyframe_n_cells(PyFrameObject* frame) {return CODE->co_nlocalsplus - CODE->co_nlocals;}
+
+        static PyObject** _pyframe_get_locals(PyFrameObject* frame) {return FRAME->localsplus;}
+        static PyObject** _pyframe_get_cells(PyFrameObject* frame) {return FRAME->localsplus + CODE->co_nlocals;}
+
+    #elif PYTELEPORT_PYTHON_VERSION == 0x030A
+        #include "tupleobject.h"
+
+        static PyObject** _pyframe_get_value_stack(PyFrameObject* frame) {return frame->f_valuestack;}
+        static int _pyframe_get_value_stack_depth(PyFrameObject* frame) {return frame->f_stackdepth;}
+
+        #define _PYFRAME_DEFINE_BLOCK_STACK_GETTER(name) static int _pyframe_get_block_stack_ ## name(PyFrameObject* frame, int i) {return frame->f_blockstack[i].name;}
+        static int _pyframe_get_block_stack_depth(PyFrameObject* frame) {return frame->f_iblock;}
+
+        static PyObject** _pyframe_get_locals(PyFrameObject* frame) {return frame->f_localsplus;}
+        static PyObject** _pyframe_get_cells(PyFrameObject* frame) {return frame->f_localsplus + frame->f_code->co_nlocals;}
+
+        static int _pyframe_n_locals(PyFrameObject* frame) {return frame->f_code->co_nlocals;}
+        static int _pyframe_n_cells(PyFrameObject* frame) {return PyTuple_Size(frame->f_code->co_freevars) + PyTuple_Size(frame->f_code->co_cellvars);}
+
+    #elif PYTELEPORT_PYTHON_VERSION == 0x0309
+
+        static PyObject** _pyframe_get_value_stack(PyFrameObject* frame) {return frame->f_valuestack;}
+        static int _pyframe_get_value_stack_depth(PyFrameObject* frame) {
+            if (frame->f_stacktop)
+                return (int) (frame->f_stacktop - frame->f_valuestack);
+            else
+                return -1;
+        }
+
+        #define _PYFRAME_DEFINE_BLOCK_STACK_GETTER(name) static int _pyframe_get_block_stack_ ## name(PyFrameObject* frame, int i) {return frame->f_blockstack[i].name;}
+        static int _pyframe_get_block_stack_depth(PyFrameObject* frame) {return frame->f_iblock;}
+
+        static PyObject** _pyframe_get_locals(PyFrameObject* frame) {return frame->f_localsplus;}
+        static PyObject** _pyframe_get_cells(PyFrameObject* frame) {return frame->f_localsplus + frame->f_code->co_nlocals;}
+
+        static int _pyframe_n_locals(PyFrameObject* frame) {return frame->f_code->co_nlocals;}
+        static int _pyframe_n_cells(PyFrameObject* frame) {return PyTuple_Size(frame->f_code->co_freevars) + PyTuple_Size(frame->f_code->co_cellvars);}
+
     #else
-      #error "PY_VERSION_HEX not defined"
+        #error "Not implemented for this cpython version"
     #endif
+
+    _PYFRAME_DEFINE_BLOCK_STACK_GETTER(b_type)
+    _PYFRAME_DEFINE_BLOCK_STACK_GETTER(b_handler)
+    _PYFRAME_DEFINE_BLOCK_STACK_GETTER(b_level)
     """
-    int _pyteleport_stackdepth(_frame* frame)
+
+    PyObject** _pyframe_get_value_stack(PyFrameObject* frame)
+    int _pyframe_get_value_stack_depth(PyFrameObject* frame)
+    int _pyframe_get_block_stack_b_type(PyFrameObject* frame, int i)
+    int _pyframe_get_block_stack_b_handler(PyFrameObject* frame, int i)
+    int _pyframe_get_block_stack_b_level(PyFrameObject* frame, int i)
+    int _pyframe_get_block_stack_depth(PyFrameObject* frame)
+
+    PyObject** _pyframe_get_locals(PyFrameObject* frame)
+    PyObject** _pyframe_get_cells(PyFrameObject* frame)
+    int _pyframe_n_locals(PyFrameObject* frame)
+    int _pyframe_n_cells(PyFrameObject* frame)
 
 
 NOTSET = object()
 
 
 cdef class FrameWrapper:
-    cdef _frame* frame
+    cdef PyFrameObject* frame
 
     def __cinit__(self, object frame):
-        self.frame = <_frame*>frame
+        self.frame = <PyFrameObject*>frame
 
     @property
     def block_stack(self):
         cdef:
             int i
-            PyTryBlock ptb
+        if _pyframe_get_block_stack_depth(self.frame) == -1:
+            raise ValueError("not implemented for this python version")
 
         result = []
-        for i in range(self.frame.f_iblock):
-            ptb = self.frame.f_blockstack[i]
-            result.append(block_stack_item(ptb.b_type, ptb.b_handler, ptb.b_level))
+        for i in range(_pyframe_get_block_stack_depth(self.frame)):
+            result.append(block_stack_item(
+                _pyframe_get_block_stack_b_type(self.frame, i),
+                _pyframe_get_block_stack_b_handler(self.frame, i),
+                _pyframe_get_block_stack_b_level(self.frame, i),
+            ))
         return result
 
     def get_value_stack(self, int stack_size = -1, object null = NULL_object):
@@ -67,38 +123,42 @@ cdef class FrameWrapper:
 
         # first, determine the stack size
         if stack_size < 0:
-            stack_size = _pyteleport_stackdepth(self.frame)  # only works for inactive generator frames
+            stack_size = _pyframe_get_value_stack_depth(self.frame)  # only works for inactive generator frames
             if stack_size < 0:
                 raise ValueError("this frame requires stack size")
 
         # second, copy stack objects
         result = []
         for i in range(stack_size):
-            stack_item = self.frame.f_valuestack[i]
+            stack_item = _pyframe_get_value_stack(self.frame)[i]
             if stack_item:
                 result.append(<object>stack_item)
             else:
                 result.append(null)
         return result
 
-    def get_locals_plus(self, object null=NULL_object):
+    def get_locals(self, object null = NULL_object):
         cdef:
-            PyObject* item
-            int i, n_locals
+            PyObject** ptr = _pyframe_get_locals(self.frame)
+            int i
 
-        code = (<object>self.frame).f_code
-        n_locals = code.co_nlocals
-        assert len(code.co_varnames) == n_locals
-        cdef:
-            int n_cells = len(code.co_cellvars)
-            int n_free = len(code.co_freevars)
-
-        locals = []
-        for i in range(n_locals + n_cells + n_free):
-            item = self.frame.f_localsplus[i]
-            if item:
-                locals.append(<object>item)
+        result = []
+        for i in range(_pyframe_n_locals(self.frame)):
+            if ptr[i]:
+                result.append(<object>ptr[i])
             else:
-                locals.append(null)
+                result.append(null)
+        return result
+
+    def get_cells(self, object null = NULL_object):
+        cdef:
+            PyObject** ptr = _pyframe_get_cells(self.frame)
+            int i
 
-        return locals[:n_locals], locals[n_locals:n_locals + n_cells], locals[n_locals + n_cells:]
+        result = []
+        for i in range(_pyframe_n_cells(self.frame)):
+            if ptr[i]:
+                result.append(<object>ptr[i])
+            else:
+                result.append(null)
+        return result
diff --git a/pyteleport/snapshot.py b/pyteleport/snapshot.py
@@ -265,10 +265,12 @@ def snapshot(topmost_frame, stack_method="predict"):
             disassemble(fs.code).print(log_bytecode)
             raise NotImplementedError(f"Cannot interpret {fs.current_opcode_repr}")
 
-        v_locals, v_cells, v_free = frame_wrapper.get_locals_plus()
-        fs = fs._replace(v_stack=vstack[:stack_size], v_locals=v_locals,
-                         v_cells=v_cells + v_free,
-                         tos_plus_one=called)
+        fs = fs._replace(
+            v_stack=vstack[:stack_size],
+            v_locals=frame_wrapper.get_locals(),
+            v_cells=frame_wrapper.get_cells(),
+            tos_plus_one=called,
+        )
 
         result.append(fs)
     logging.debug("  verifying frame stack continuity ...")

diff --git a/pyteleport/tests/test_frame.py b/pyteleport/tests/test_frame.py
@@ -8,11 +8,12 @@ def test_basic():
     wrapper = FrameWrapper(frame)
     assert wrapper.block_stack == []
     assert wrapper.get_value_stack(0) == []
-    locals_, cells, free = wrapper.get_locals_plus()
+    locals_ = wrapper.get_locals()
+    cells = wrapper.get_cells()
     for obj in frame, wrapper:
         assert obj in locals_
     assert cells == []
-    assert free == []
+    assert cells == []
 
 
 def test_value_stack():
@@ -53,7 +54,7 @@ def generator():
     assert wrapper.block_stack == []
     range_itr, = wrapper.get_value_stack()
     assert "range_iterator" in str(range_itr)
-    locals_, cells, free = wrapper.get_locals_plus()
+    locals_ = wrapper.get_locals()
+    cells = wrapper.get_cells()
     assert locals_ == [0]
     assert cells == []
-    assert free == []