PyO3 / PyO3/pyo3

Implement a safe API wrapping PyEval_SetProfile

Open
#4,008 9 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Rust
Stars
16.2k
Forks
1k
Avg merge
2d 6h
Merged PRs (30d)
66

Description

PyEval_SetProfile is the C-api equivalent to sys.set_profile. Using PyEval_SetProfile from Rust is preferable to using sys.set_profile because it avoids overhead, but currently requires using unsafe. I would like to add a new safe api to replace code like the following:

/// Wrap pyo3-ffi/src/cpython/pystate.rs#L18-L25
enum Event {
    Call,
    Exception,
    Line,
    Return,
    CCall,
    CException,
    CReturn,
    Opcode,
}

impl Event {
    fn from_c(what: c_int) -> Self {
         match what {
             PyTrace_CALL => Self::Call,
             PyTrace_EXCEPTION => Self::Exception,
             PyTrace_LINE => Self::Line,
             PyTrace_RETURN => Self::Return,
             PyTrace_C_CALL => Self::CCall,
             PyTrace_C_EXCEPTION => Self::CException,
             PyTrace_C_RETURN => Self::CReturn,
             PyTrace_OPCODE => Self::Opcode,
         }
    }
}

#[pyclass]
struct Profiler {
    // Useful fields
}

impl Profiler {
    fn profile(
        &mut self,
        frame: PyObject,
        arg: Option<PyObject>,
        event: Event,
        py: Python,
    ) -> PyResult<()> {
        // Custom profiling logic
        Ok(())
    }
}

pub extern "C" fn profile_callback(
    _obj: *mut ffi::PyObject,
    _frame: *mut ffi::PyFrameObject,
    what: c_int,
    _arg: *mut ffi::PyObject,
) -> c_int {
    let event = Event::from_c(what);
    // An optimisation for my use case that might not be worth trying to allow upstream
    // match event {
    //     Event::Call => (),
    //     Event::Return => (),
    //     _ => return 0;
    //}
    let _frame = _frame as *mut ffi::PyObject;
    Python::with_gil(|py| {
        // Safety:
        //
        // `from_borrowed_ptr_or_err` must be called in an unsafe block.
        //
        // `_obj` is a reference to our `Profiler` wrapped up in a Python object, so
        // we can safely convert it from an `ffi::PyObject` to a `PyObject`.
        //
        // We borrow the object so we don't break reference counting.
        //
        // https://docs.rs/pyo3/latest/pyo3/struct.Py.html#method.from_borrowed_ptr_or_err
        // https://docs.python.org/3/c-api/init.html#c.Py_tracefunc
        let obj = match unsafe { PyObject::from_borrowed_ptr_or_err(py, _obj) } {
            Ok(obj) => obj,
            Err(err) => {
                err.restore(py);
                return -1;
            }
        };
        let mut profiler = match obj.extract::<PyRefMut<Profiler>>(py) {
            Ok(profiler) => profiler,
            Err(err) => {
                err.restore(py);
                return -1;
            }
        };

        // Safety:
        //
        // `from_borrowed_ptr_or_err` must be called in an unsafe block.
        //
        // `_frame` is an `ffi::PyFrameObject` which can be converted safely
        // to a `PyObject`. We can later convert it into a `pyo3::types::PyFrame`.
        //
        // We borrow the object so we don't break reference counting.
        //
        // https://docs.rs/pyo3/latest/pyo3/struct.Py.html#method.from_borrowed_ptr_or_err
        // https://docs.python.org/3/c-api/init.html#c.Py_tracefunc
        let frame = match unsafe { PyObject::from_borrowed_ptr_or_err(py, _frame) } {
            Ok(frame) => frame,
            Err(err) => {
                err.restore(py);
                return -1;
            }
        };

        // Safety:
        //
        // `from_borrowed_ptr_or_opt` must be called in an unsafe block.
        //
        // `_arg` is either a `Py_None` (PyTrace_CALL) or any PyObject (PyTrace_RETURN) or
        // NULL (PyTrace_RETURN).
        //
        // We borrow the object so we don't break reference counting.
        //
        // https://docs.rs/pyo3/latest/pyo3/struct.Py.html#method.from_borrowed_ptr_or_opt
        // https://docs.python.org/3/c-api/init.html#c.Py_tracefunc
        let arg = unsafe { PyObject::from_borrowed_ptr_or_opt(py, _arg) };
        // `_arg` is `NULL` when the frame exits with an exception unwinding instead of a normal return.
        // So it might be possible to make `arg` a `PyResult` here instead of an option, but I haven't worked out the detail of how that would work. 

        match profiler.profile(frame, arg, event, py) {
            Ok(_) => 0,
            Err(err) => {
                err.restore(py);
                return -1;
            }
        }
    })

}

#[pyfunction]
fn register_profiler() -> PyResult<()> {
    Python::with_gil(|py| {
        let profiler = Profiler::new();
        unsafe {
            ffi::PyEval_SetProfile(Some(profile_callback), profiler.into_ptr());
        }
    }
}

Using the safe API could look something like:

#[pyclass]
struct Profiler {
    // Useful fields
}

impl Profiler {
    fn profile(
        &mut self,
        frame: PyObject,
        arg: Option<PyObject>,
        event: pyo3::introspection::Event,
        py: Python,
    ) -> PyResult<()> {
        // Custom profiling logic
        Ok(())
    }
}

#[pyfunction]
fn register_profiler() -> PyResult<()> {
    Python::with_gil(|py| {
        let profiler = Profiler::new();
        pyo3::introspection::set_profile(profiler.profile, profiler)?;
        Ok(())
    }
}

In the C api it is also possible to omit the Profiler struct and just pass a profiling callback. It would be nice to support this too, but I'm not sure there's a clean way without a separate function:

fn profile(
    frame: PyObject,
    arg: Option<PyObject>,
    event: pyo3::introspection::Event,
    py: Python,
) -> PyResult<()> {
    // Custom profiling logic
    Ok(())
}

#[pyfunction]
fn register_profiler() -> PyResult<()> {
    Python::with_gil(|py| {
        pyo3::introspection::set_profile2(profile)?;
        Ok(())
    }
}

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading pyo3-ffi/src/cpython/pystate.rs around the PyTrace event definitions and the PyEval_SetProfile C API documentation. Compare the proposed profiler and callback forms, including their event, frame, argument, and error-handling behavior. Done means a safe public API covers the requested profiling use cases without requiring callers to write the shown unsafe callback wrapper.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, rust
Domain
api, developer-experience
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.