Skip to main content

comfy_builder_core/
node.rs

1use pyo3::types::{PyCFunction, PyDict, PyList, PyTuple};
2use pyo3::{Bound, PyAny, PyErr, PyResult, Python};
3use std::error::Error;
4use std::ops::Deref;
5
6pub trait Node: Default {
7    type In: In;
8    type Out: Out;
9
10    type Error: Into<Box<dyn Error + Send + Sync>> + 'static;
11
12    fn new() -> Self {
13        Default::default()
14    }
15
16    fn initialize_inputs<'py>(
17        &self,
18        kwargs: Kwargs<'py>,
19    ) -> Result<Self::In, <Self::In as TryFrom<Kwargs<'py>>>::Error> {
20        Self::In::try_from(kwargs)
21    }
22
23    fn execute(&self, input: Self::In) -> Result<Self::Out, Self::Error>;
24}
25
26pub trait NodeFunctionProvider {
27    fn define_fn(python: Python) -> PyResult<Bound<PyCFunction>>;
28    fn execute_fn(python: Python) -> PyResult<Bound<PyCFunction>>;
29}
30
31pub trait In: for<'py> TryFrom<Kwargs<'py>> {
32    fn blueprints<'py>(python: Python<'py>, io: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyList>>;
33    fn is_list() -> bool;
34}
35
36pub trait Out {
37    fn blueprints<'py>(python: Python<'py>, io: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyList>>;
38    fn to_schema(self, python: Python) -> PyResult<Bound<PyTuple>>;
39}
40
41pub struct Kwargs<'py>(pub Option<Bound<'py, PyDict>>);
42
43impl<'py> From<Option<Bound<'py, PyDict>>> for Kwargs<'py> {
44    fn from(value: Option<Bound<'py, PyDict>>) -> Self {
45        Kwargs(value)
46    }
47}
48
49impl<'a> Deref for Kwargs<'a> {
50    type Target = Option<Bound<'a, PyDict>>;
51
52    fn deref(&self) -> &Self::Target {
53        &self.0
54    }
55}
56
57//--- Allow Input / Output to be set as empty unit type
58
59impl Out for () {
60    fn blueprints<'py>(python: Python<'py>, _: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyList>> {
61        Ok(PyList::empty(python))
62    }
63
64    fn to_schema(self, python: Python) -> PyResult<Bound<PyTuple>> {
65        Ok(PyTuple::empty(python))
66    }
67}
68
69impl<'py> TryFrom<Kwargs<'py>> for () {
70    type Error = PyErr;
71
72    fn try_from(_: Kwargs) -> Result<Self, Self::Error> {
73        Ok(())
74    }
75}
76
77impl In for () {
78    fn blueprints<'py>(python: Python<'py>, _: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyList>> {
79        Ok(PyList::empty(python))
80    }
81
82    fn is_list() -> bool {
83        false
84    }
85}