comfy_builder_core/
registry.rs

1use crate::node::NodeFunctionProvider;
2use pyo3::prelude::PyAnyMethods;
3use pyo3::types::{PyCFunction, PyDict, PyDictMethods};
4use pyo3::{Bound, PyAny, PyResult, Python};
5
6type FactoryFn = for<'py> fn(python: Python<'py>) -> PyResult<(Bound<'py, PyCFunction>, Bound<'py, PyCFunction>)>;
7
8#[derive(Debug)]
9pub struct NodeRegistration {
10    factory: FactoryFn,
11}
12
13#[derive(Debug)]
14pub struct EnumRegistration {
15    pub name: &'static str,
16}
17
18impl NodeRegistration {
19    pub const fn new<T: NodeFunctionProvider>() -> Self {
20        Self {
21            factory: |python| Ok((T::define_fn(python)?, T::execute_fn(python)?)),
22        }
23    }
24
25    pub fn create_node<'a, 'py>(
26        &self,
27        python: Python<'py>,
28        decorator: &'a Bound<'py, PyAny>,
29        type_fn: &'a Bound<'py, PyAny>,
30        comfy_node: &'a Bound<'py, PyAny>,
31        module_name: &'static str,
32    ) -> PyResult<Bound<'py, PyAny>> {
33        let (define, execute) = (self.factory)(python)?;
34        let methods = PyDict::new(python);
35
36        methods.set_item("define_schema", decorator.call1((define,))?)?;
37        methods.set_item("execute", decorator.call1((execute,))?)?;
38
39        type_fn.call1((format!("{}_node", module_name), (comfy_node,), methods))
40    }
41}
42
43inventory::collect!(NodeRegistration);
44inventory::collect!(EnumRegistration);
comfy_builder_core/registry.rs

comfy_builder_core/
registry.rs
