comfy_builder_core/
registry.rs1use 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);