comfy_builder_core/
registry.rs

1use pyo3::prelude::PyModule;
2use pyo3::types::{PyDict, PyDictMethods, PyModuleMethods};
3use pyo3::{Bound, PyClass, PyResult, Python};
4
5pub trait Registerable: PyClass {}
6
7#[derive(Debug)]
8pub struct NodeRegistration {
9    inner: fn(
10        python: Python,
11        module: &Bound<PyModule>,
12        class: &Bound<PyDict>,
13        display: &Bound<PyDict>,
14    ) -> PyResult<()>,
15}
16
17impl NodeRegistration {
18    pub const fn new<T: Registerable>() -> Self {
19        Self {
20            inner: |python, module, class, display| {
21                module.add_class::<T>()?;
22
23                let type_name = std::any::type_name::<T>();
24                let class_name = type_name.split("::").last().unwrap_or(type_name);
25
26                class.set_item(class_name, python.get_type::<T>())?;
27                display.set_item(class_name, class_name)?;
28
29                Ok(())
30            },
31        }
32    }
33
34    pub fn register(
35        &self,
36        python: Python,
37        module: &Bound<PyModule>,
38        class: &Bound<PyDict>,
39        display: &Bound<PyDict>,
40    ) -> PyResult<()> {
41        (self.inner)(python, module, class, display)
42    }
43}
44
45inventory::collect!(NodeRegistration);