Adding pyo3 to td-rs for Better Python Integration
pyo3 and td-rs
TouchDesigner is, in many ways, a big Python appilcation. While lots can be done with the node graph, advanced users will often find themselves writing loads of Python. Our initial Python integration was functional but left much to be desired, requiring writing Python extensions in the style of the Python C API rather than idiomatic Rust.
This post details the work integrating pyo3, an incredible Rust-Python interoperability library, with td-rs. It might be bold to claim, but the experience of creating Python-enabled TouchDesigner operators in Rust is probably the strongest feature of td-rs. However, unlike most pyo3 use cases, TouchDesigner's architecture presents unique challenges that required customizing pyo3's memory model and doing other forms of dark magic.
python in touchdesigner
It's important to understand how TouchDesigner exposes C++ objects to Python. In TouchDesigner, all operators are C++ classes that inherit from a base operator class. When you access an operator from Python, TouchDesigner creates a Python object wrapper that delegates to the C++ operator instance via any methods or properties defined via the C Python API.
TouchDesigner's Python integration involves embedding Python objects at specific memory offsets within C++ operator classes. Each operator class has a PY_Struct field that provides access to the Python interpreter and allows TouchDesigner to expose the operator's functionality to Python.
This is fundamentally different from how most Rust-Python integrations work, where Rust code typically creates and manages the lifecycle of Python objects. In our case, TouchDesigner owns the objects, and we need to hook into their existing lifecycle.
previous python integration approach
Our initial approach used a combination of custom derive macros and direct use of the Python C API to expose Rust operators to Python. This was meant to simulate the convenience of the pyo3 proc macros, but ultimately had much worse ergonomics. Here's an example of how we previously exposed a Python-enabled CHOP operator:
#[derive(PyOp, Debug)]
pub struct PythonChop {
#[py(doc = "Get or Set the speed modulation.", auto_cook)]
speed: f32,
#[py(get, doc = "Get executed count.")]
execute_count: u32,
offset: f32,
}
#[py_op_methods]
impl PythonChop {
fn reset_filter(&mut self) {
self.offset = 0.0;
}
#[py_meth]
pub unsafe fn reset(
&mut self,
_args: *mut *mut pyo3_ffi::PyObject,
_nargs: usize,
) -> *mut pyo3_ffi::PyObject {
self.reset_filter();
let none = pyo3_ffi::Py_None();
pyo3_ffi::Py_INCREF(none);
none
}
}
This approach worked, but had several drawbacks:
- ~ It required heavy use of
unsafecode to interact with the Python C API - ~ Developers had to manually manage Python reference counting, similarly to how Python interop works in TouchDesigner's C++ API examples
- ~ Type conversions were manual and error-prone, including platform specific differences about how CPython handles numeric types
- ~ Error handling was minimal, often resulting in crashes rather than proper Python exceptions
Under the hood, our macros were generating significant amounts of FFI code. For example, the PyOp derive macro would generate getter and setter methods for each field marked with #[py], and the py_op_methods macro would generate the necessary C FFI shims for each method. This worked, but it was brittle and barely better than writing all the interop yourself.
pyo3
pyo3 is a gorgeous library that does lots of proc macro based metaprogramming to make interop with Python totally effortless. It provides high-level, safe Rust bindings to the Python C API, handling all the low-level details like reference counting, type conversions, and error handling.
With pyo3, Python integration in Rust looks like this:
use pyo3::prelude::*;
#[pyclass]
struct MyClass {
#[pyo3(get, set)]
value: i32,
}
#[pymethods]
impl MyClass {
#[new]
fn new(value: i32) -> Self {
Self { value }
}
fn double(&self) -> i32 {
self.value * 2
}
}
This is much cleaner than working directly with the Python C API, as pyo3 handles all the low-level details for you. However, pyo3 assumes it's creating and managing the lifecycle of the Python objects, particularly Python classes. In TouchDesigner, the objects are created and managed by the plugin system, so we need to customize pyo3 to work within TouchDesigner's memory model.
memory model challenges
In standard pyo3 usage, the library creates a Python object that wraps a Rust struct, and the Rust struct's memory layout is controlled by pyo3. TouchDesigner, however, already has its own Python object system and memory layout, with Python object data at specific offsets within C++ operator structs.
This presents several challenges:
- ~ pyo3 normally calculates field offsets relative to the start of the Python object, but in TouchDesigner, the Python object is embedded within the larger operator structure
- ~ pyo3 assumes it controls the lifecycle of the Python object, but in TouchDesigner, the lifecycle is controlled by the plugin system
- ~ pyo3's type conversion system assumes it's working with objects it created, not with pre-existing objects created by TouchDesigner
The most significant issue was with field access. In pyo3, when a Python script accesses a field of a Rust struct, pyo3 calculates the field's offset from the start of the struct and reads or writes the value directly. This works fine when pyo3 has created the object, but in TouchDesigner, the struct is part of a larger C++ object, and the field offset calculation is more complex.
Let's look at the problem in pyo3's field access code:
// From pyo3's impl_/pyclass.rs
fn field_from_object<ClassT, FieldT, Offset>(obj: *mut ffi::PyObject) -> *mut FieldT
where
ClassT: PyClass,
Offset: OffsetCalculator<ClassT, FieldT>,
{
unsafe { obj.cast::<u8>().add(Offset::offset()).cast::<FieldT>() }
}
This function takes a PyObject pointer, casts it to a byte pointer, adds the offset of the field, and then casts it to a pointer to the field type. This assumes that the PyObject pointer points to the start of the Rust struct, which is not the case in TouchDesigner.
customizing pyo3's memory model
Our solution involved the following:
- ~ Forking pyo3 and modifying its memory model to support TouchDesigner's object layout
- ~ Creating specialized traits for extracting Rust references from TouchDesigner's Python objects
- ~ Implementing a bridge between TouchDesigner's Python object system and pyo3's type system
- ~ Updating our Python integration macros to use pyo3's traits and attributes
extracting py classes
First, we needed to modify pyo3's memory model to work with TouchDesigner's object layout. We made several key changes:
- ~ Added two new traits,
ExtractPyClassRefandExtractPyClassRefMut, which allow customizing how Rust references are extracted from Python objects - ~ Modified the field access functions to use our custom traits instead of direct offset calculations
- ~ Used the nightly
min-specializationfeature, which allows us to provide custom implementations of these traits for TouchDesigner operators
Here's a simplified version of our changes to the field access code:
// Old field access function in pyo3
fn field_from_object<ClassT, FieldT, Offset>(obj: *mut ffi::PyObject) -> *mut FieldT
where
ClassT: PyClass,
Offset: OffsetCalculator<ClassT, FieldT>,
{
unsafe { obj.cast::<u8>().add(Offset::offset()).cast::<FieldT>() }
}
// New field access function in our modified pyo3
fn field_from_object<ClassT, FieldT, Offset>(clazz: &ClassT) -> *const FieldT
where
ClassT: PyClass + for<'a, 'py> ExtractPyClassRef<'a, 'py>,
Offset: OffsetCalculator<ClassT, FieldT>,
{
unsafe { (clazz as *const ClassT).cast::<u8>().add(Offset::offset()).cast::<FieldT>() }
}
The key difference is that the new function takes a reference to the Rust struct directly, rather than a PyObject pointer. This allows us to calculate field offsets correctly for TouchDesigner's object layout.
custom extraction traits
Next, we created new traits for extracting Rust references from TouchDesigner's Python objects:
pub trait ExtractPyClassRef<'a, 'py>
where
Self: PyClass,
{
fn extract_ref(
obj: &'a Bound<'py, PyAny>,
holder: &'a mut Option<PyRef<'py, Self>>,
) -> PyResult<&'a Self>;
}
pub trait ExtractPyClassRefMut<'a, 'py>
where
Self: PyClass<Frozen = False>,
{
fn extract_mut(
obj: &'a Bound<'py, PyAny>,
holder: &'a mut Option<PyRefMut<'py, Self>>,
) -> PyResult<&'a mut Self>;
}
These traits define how to extract a reference to a Rust struct from a Python object. They take a Python object reference (Bound<'py, PyAny>) and a holder for the extracted reference, and return a result containing a reference to the Rust struct.
We then provided default implementations of these traits using unstablespecialization:
impl<'a, 'py, T> ExtractPyClassRef<'a, 'py> for T
where
T: PyClass,
{
default fn extract_ref(
obj: &'a Bound<'py, PyAny>,
holder: &'a mut Option<PyRef<'py, T>>,
) -> PyResult<&'a T> {
Ok(&*holder.insert(obj.extract()?))
}
}
impl<'a, 'py, T> ExtractPyClassRefMut<'a, 'py> for T
where
T: PyClass<Frozen = False>,
{
default fn extract_mut(
obj: &'a Bound<'py, PyAny>,
holder: &'a mut Option<PyRefMut<'py, T>>,
) -> PyResult<&'a mut T> {
Ok(&mut *holder.insert(obj.extract()?))
}
}
These default implementations provide the standard pyo3 behavior for normal Rust types. For TouchDesigner operators, we specialize these traits to extract references from TouchDesigner's Python objects. Specialization is unstable but is incredibly useful for these kinds of scenarios where we want to provide a blanket implementation that we override in our downstream library. This technique is also used in td-rs so that users don't have to implement every possible trait we define.
bridging touchdesigner and pyo3
The most crucial part of our solution is the implementation of these traits for TouchDesigner operators. Here's how we extract a reference to a TouchDesigner operator from a Python object via cxx:
impl<'a, 'py> ExtractPyClassRef<'a, 'py> for PythonChop {
fn extract_ref(
obj: &'a Bound<'py, PyAny>,
holder: &'a mut Option<PyRef<'py, Self>>,
) -> PyResult<&'a Self> {
unsafe {
// Get the TouchDesigner Python structure
let me = obj.as_ptr();
let py_struct = me as *mut cxx::PY_Struct;
let info = cxx::PY_GetInfo {
autoCook: true,
reserved: [0; 50],
};
// Access the Python context
// Pin is safe here because we're operating in the context of cook where
// TouchDesigner won't move out from under us until the execute lifecycle ends
// We should probably be better about writing explicit safety comments (:
let mut ctx = Pin::new_unchecked(&mut *cxx::getPyContext(py_struct));
// Get our operator instance via TouchDesigner's API
let me = ctx.getNodeInstance(&info, std::ptr::null_mut());
if me.is_null() {
return Err(PyTypeError::new_err("operator is null"));
}
// Cast to our Rust type via some sketchy cxx on the C++ side
let py_op = {
let me = cxx::plugin_cast(me);
// Unwrap the cxx delegate and finally return our actual instance
let me = me.as_plugin().inner();
&*(me as *const PythonChop)
};
Ok(py_op)
}
}
}
This code navigates TouchDesigner's object hierarchy to extract a reference to our Rust struct. It first gets the PY_Struct from the Python object, then uses TouchDesigner's API to get a pointer to the operator instance, and finally casts this pointer to our Rust type.
We also implemented ExtractPyClassRefMut with the following key difference:
Pin::new_unchecked(&mut *py_ctx).makeNodeDirty(std::ptr::null_mut());
We marks the node as dirty (causing it to be re-cooked on the next frame) and returns a mutable reference to our Rust struct. It's also all a little sketchy!
python integration macros
With these changes in place, we updated our PyOp derive macro to use pyo3's attributes and traits:
#[proc_macro_derive(PyOp, attributes(pyo3))]
pub fn derive_py_op(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as DeriveInput);
impl_py_op(&input)
}
Now, instead of generating custom getter and setter methods for each field, the macro implements the necessary pyo3 traits and hooks into our custom extraction mechanism.
The result is a much cleaner API for creating Python-enabled TouchDesigner operators:
#[derive(PyOp)]
#[pyclass(unsendable)]
pub struct PythonChop {
info: NodeInfo,
#[pyo3(get, set)]
speed: f32,
#[pyo3(get)]
execute_count: u32,
offset: f32,
params: PythonChopParams,
}
#[pymethods]
impl PythonChop {
pub fn reset(&mut self) {
self.offset = 0.0;
}
}
This looks much more like standard pyo3 code, with the addition of our PyOp derive macro that hooks into TouchDesigner's Python object system.
improved python integration
Our pyo3 integration significantly improves the developer experience in several ways:
method definitions
Previously, creating a Python-exposed method required manually implementing FFI shims and dealing with raw pointers:
#[py_meth]
pub unsafe fn reset(
&mut self,
_args: *mut *mut pyo3_ffi::PyObject,
_nargs: usize,
) -> *mut pyo3_ffi::PyObject {
self.reset_filter();
let none = pyo3_ffi::Py_None();
pyo3_ffi::Py_INCREF(none);
none
}
Now, we can use pyo3's #[pymethods] attribute and write normal Rust methods:
#[pymethods]
impl PythonChop {
pub fn reset(&mut self) {
self.offset = 0.0;
}
}
pyo3 handles all the FFI details, including type conversions, reference counting, and error handling.
field access
Field access is similarly improved. Previously, we had to manually implement getters and setters for each field. While code superficially looked the same, it would generate brittle FFI methods to access the fields:
#[py(get, set, doc = "Get or Set the speed modulation.", auto_cook)]
speed: f32,
Now, we use pyo3's built-in field attributes which are significantly robust and calculate the field offsets for us:
#[pyo3(get, set)]
speed: f32,
pyo3 handles all the details of exposing fields to Python, including type conversions and error handling.
calling python from rust
Calling Python functions from Rust is also significantly better. Previously, we had to use the CPython directly:
let arg_tuple = self.info.context().create_arguments_tuple(1);
unsafe {
pyo3_ffi::PyTuple_SET_ITEM(
arg_tuple,
1,
pyo3_ffi::PyFloat_FromDouble(self.params.speed as std::ffi::c_double),
);
let res = self.info.context().call_python_callback(
"getSpeedAdjust",
arg_tuple,
std::ptr::null_mut(),
);
if !res.is_null() {
if pyo3_ffi::PyFloat_Check(res) != 0 {
self.params.speed = pyo3_ffi::PyFloat_AsDouble(res) as f32;
}
pyo3_ffi::Py_DECREF(res);
}
}
Now, we can use pyo3's higher-level API:
Python::with_gil(|py| {
self.info.context().call_python_callback(
py,
"getSpeedAdjust",
(self.speed,),
None,
|py, res| {
if let Ok(speed) = res.extract::<f32>(py) {
self.params.speed *= speed;
}
},
)
})
.unwrap();
This is not only shorter but also safer, with proper GIL management, type-safe tuples for arguments, and proper error handling.
handling python getsets and methods
One of the more complex aspects of our integration was handling Python's GetSet protocol and method definitions. In TouchDesigner, when you define Python properties or methods on an operator, it needs to register these with the Python C API.
Previously, we manually generated the necessary PyGetSetDef and PyMethodDef structures for each property and method. With pyo3, we leverage its existing machinery for generating these structures, but we need to extract them from pyo3's internal representation.
For methods, we do this by accessing pyo3's PyClassImplCollector to get the method definitions. Note, this is totally not their intention (see the impl_ namespace) and is likely to break in future versions! Sometimes you have to live on the edge to get the API you want, though:
impl PyMethods for PythonChop {
fn get_methods() -> &'static [pyo3::ffi::PyMethodDef] {
let clazz = pyo3::impl_::pyclass::PyClassImplCollector::<PythonChop>::new();
let methods = clazz.py_methods();
let mut method_defs = Vec::new();
for method in methods.methods {
let method_def = match method {
pyo3::impl_::pyclass::MaybeRuntimePyMethodDef::Runtime(m) => &m(),
pyo3::impl_::pyclass::MaybeRuntimePyMethodDef::Static(m) => m,
};
match method_def {
pyo3::PyMethodDefType::Method(m) => {
method_defs.push(m.as_method_def());
}
_ => {}
}
}
method_defs.leak()
}
}
This code extracts the method definitions from pyo3's internal representation and converts them to the format expected by TouchDesigner.
Similarly, for getters and setters, we extract the definitions from pyo3's internal representation:
impl PyGetSets for PythonChop {
fn get_get_sets() -> &'static [pyo3::ffi::PyGetSetDef] {
let clazz = pyo3::impl_::pyclass::PyClassImplCollector::<PythonChop>::new();
let methods = clazz.py_methods();
let mut getset_builders = std::collections::HashMap::<&std::ffi::CStr, pyo3::pyclass::create_type_object::GetSetDefBuilder>::new();
for method in methods.methods {
let method_def = match method {
pyo3::impl_::pyclass::MaybeRuntimePyMethodDef::Runtime(m) => &m(),
pyo3::impl_::pyclass::MaybeRuntimePyMethodDef::Static(m) => m,
};
match method_def {
pyo3::PyMethodDefType::Getter(getter) => {
getset_builders
.entry(getter.name)
.or_default()
.add_getter(getter)
}
pyo3::PyMethodDefType::Setter(setter) => {
getset_builders
.entry(setter.name)
.or_default()
.add_setter(setter)
}
_ => {}
}
}
// Process additional inherited methods from parent classes
let items = PythonChop::items_iter();
for item in items {
for method in item.methods {
let method_def = match method {
pyo3::impl_::pyclass::MaybeRuntimePyMethodDef::Runtime(m) => &m(),
pyo3::impl_::pyclass::MaybeRuntimePyMethodDef::Static(m) => m,
};
match method_def {
pyo3::PyMethodDefType::Getter(getter) => {
getset_builders
.entry(getter.name)
.or_default()
.add_getter(getter)
}
pyo3::PyMethodDefType::Setter(setter) => {
getset_builders
.entry(setter.name)
.or_default()
.add_setter(setter)
}
_ => {}
}
}
}
let mut getset_destructors = Vec::with_capacity(getset_builders.len());
let property_defs: Vec<pyo3::ffi::PyGetSetDef> = getset_builders
.iter()
.map(|(name, builder)| {
let (def, destructor) = builder.as_get_set_def(name);
getset_destructors.push(destructor);
def
})
.collect();
// We just have to leak these to keep them alive
getset_destructors.leak();
property_defs.leak()
}
}
This code extracts the getter and setter definitions from pyo3's internal representation and converts them to the FFI format expected by TouchDesigner. We need to handle both methods defined directly on the class and methods inherited from parent classes. Some small additional changes were necessary here in order to support this.
conclusion
Integrating pyo3 with TouchDesigner's Python object system was a complex task that required modifying pyo3's memory model and internal implementation details. However, the result is a much better API creating Python-enabled TouchDesigner operators in Rust.
These changes allow us to use pyo3's high-level API for Python integration while still working within TouchDesigner's object model. The result is a much more pleasant developer experience, with less boilerplate, better type safety, and improved error handling.
This is part of our ongoing mission to make td-rs a first-class option for developing TouchDesigner plugins in Rust, combining the safety and performance of Rust with the flexibility and ease of use that TouchDesigner users expect.
It would be awesome if pyo3 could in the future include some of this machinery to make these kinds of integrations easier!