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

comfy_builder_core/
registry.rs
