test: add CodeObjectWrapper tests

Author: tzanko-matev

.agents/tasks/2025/08/21-0939-codetype-interface

@@ -33,4 +33,6 @@ According to the PyO3 documentation it is preferred to use `Bound<'_, T>` instea
 * https://pyo3.rs/v0.25.1/types.html
 * https://docs.rs/pyo3/0.25.1/pyo3/index.html
 
-Also please add usage examples to the design documentation
+Also please add usage examples to the design documentation
+--- FOLLOW UP TASK ---
+Implement the CodeObjectWrapper as designed. Update the Tracer trait as well as the callback_xxx functions accordingly. Write a comprehensive unit tests for CodeObjectWrapper.

codetracer-python-recorder/Cargo.lock

@@ -18,6 +18,7 @@ default = ["extension-module"]
 pyo3 = { version = "0.25.1" }
 runtime_tracing = "0.14.0"
 bitflags = "2.4"
+once_cell = "1.19"
 
 [dev-dependencies]
 pyo3 = { version = "0.25.1", features = ["auto-initialize"] }

codetracer-python-recorder/Cargo.toml

@@ -0,0 +1,130 @@
+use once_cell::sync::OnceCell;
+use pyo3::prelude::*;
+use pyo3::types::PyCode;
+
+/// A wrapper around Python `code` objects providing cached access to
+/// common attributes and line information.
+pub struct CodeObjectWrapper {
+    obj: Py<PyCode>,
+    id: usize,
+    cache: CodeObjectCache,
+}
+
+#[derive(Default)]
+struct CodeObjectCache {
+    filename: OnceCell<String>,
+    qualname: OnceCell<String>,
+    firstlineno: OnceCell<u32>,
+    argcount: OnceCell<u16>,
+    flags: OnceCell<u32>,
+    lines: OnceCell<Vec<LineEntry>>,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct LineEntry {
+    pub offset: u32,
+    pub line: u32,
+}
+
+impl CodeObjectWrapper {
+    /// Construct from a `CodeType` object. Computes `id` eagerly.
+    pub fn new(_py: Python<'_>, obj: &Bound<'_, PyCode>) -> Self {
+        let id = obj.as_ptr() as usize;
+        Self {
+            obj: obj.clone().unbind(),
+            id,
+            cache: CodeObjectCache::default(),
+        }
+    }
+
+    /// Borrow the owned `Py<PyCode>` as a `Bound<'py, PyCode>`.
+    pub fn as_bound<'py>(&'py self, py: Python<'py>) -> Bound<'py, PyCode> {
+        self.obj.bind(py).clone()
+    }
+
+    /// Return the stable identity of the code object (equivalent to `id(code)`).
+    pub fn id(&self) -> usize {
+        self.id
+    }
+
+    pub fn filename<'py>(&'py self, py: Python<'py>) -> PyResult<&'py str> {
+        let value = self.cache.filename.get_or_try_init(|| -> PyResult<String> {
+            let s: String = self
+                .as_bound(py)
+                .getattr("co_filename")?
+                .extract()?;
+            Ok(s)
+        })?;
+        Ok(value.as_str())
+    }
+
+    pub fn qualname<'py>(&'py self, py: Python<'py>) -> PyResult<&'py str> {
+        let value = self.cache.qualname.get_or_try_init(|| -> PyResult<String> {
+            let s: String = self
+                .as_bound(py)
+                .getattr("co_qualname")?
+                .extract()?;
+            Ok(s)
+        })?;
+        Ok(value.as_str())
+    }
+
+    pub fn first_line(&self, py: Python<'_>) -> PyResult<u32> {
+        let value = *self.cache.firstlineno.get_or_try_init(|| -> PyResult<u32> {
+            let v: u32 = self
+                .as_bound(py)
+                .getattr("co_firstlineno")?
+                .extract()?;
+            Ok(v)
+        })?;
+        Ok(value)
+    }
+
+    pub fn arg_count(&self, py: Python<'_>) -> PyResult<u16> {
+        let value = *self.cache.argcount.get_or_try_init(|| -> PyResult<u16> {
+            let v: u16 = self
+                .as_bound(py)
+                .getattr("co_argcount")?
+                .extract()?;
+            Ok(v)
+        })?;
+        Ok(value)
+    }
+
+    pub fn flags(&self, py: Python<'_>) -> PyResult<u32> {
+        let value = *self.cache.flags.get_or_try_init(|| -> PyResult<u32> {
+            let v: u32 = self
+                .as_bound(py)
+                .getattr("co_flags")?
+                .extract()?;
+            Ok(v)
+        })?;
+        Ok(value)
+    }
+
+    fn lines<'py>(&'py self, py: Python<'py>) -> PyResult<&'py [LineEntry]> {
+        let vec = self.cache.lines.get_or_try_init(|| -> PyResult<Vec<LineEntry>> {
+            let mut entries = Vec::new();
+            let iter = self.as_bound(py).call_method0("co_lines")?;
+            let iter = iter.try_iter()?;
+            for item in iter {
+                let (start, _end, line): (u32, u32, Option<u32>) = item?.extract()?;
+                if let Some(line) = line {
+                    entries.push(LineEntry { offset: start, line });
+                }
+            }
+            Ok(entries)
+        })?;
+        Ok(vec.as_slice())
+    }
+
+    /// Return the source line for a given instruction offset using a binary search.
+    pub fn line_for_offset(&self, py: Python<'_>, offset: u32) -> PyResult<Option<u32>> {
+        let lines = self.lines(py)?;
+        match lines.binary_search_by_key(&offset, |e| e.offset) {
+            Ok(idx) => Ok(Some(lines[idx].line)),
+            Err(0) => Ok(None),
+            Err(idx) => Ok(Some(lines[idx - 1].line)),
+        }
+    }
+}

codetracer-python-recorder/src/code_object.rs

@@ -3,7 +3,9 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 
+pub mod code_object;
 pub mod tracer;
+pub use crate::code_object::CodeObjectWrapper;
 pub use crate::tracer::{install_tracer, uninstall_tracer, EventSet, Tracer};
 
 /// Global flag tracking whether tracing is active.