dora/apis/python/node/src/lib.rs

#![allow(clippy::borrow_deref_ref)] // clippy warns about code generated by #[pymethods]

use arrow::pyarrow::{FromPyArrow, ToPyArrow};
use dora_node_api::merged::{MergeExternalSend, MergedEvent};
use dora_node_api::{DoraNode, EventStream};
use dora_operator_api_python::{pydict_to_metadata, PyEvent};
use dora_ros2_bridge_python::Ros2Subscription;
use eyre::Context;
use futures::{Stream, StreamExt};
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict};

/// The custom node API lets you integrate `dora` into your application.
/// It allows you to retrieve input and send output in any fashion you want.
///
/// Use with:
///
/// ```python
/// from dora import Node
///
/// node = Node()
/// ```
///
#[pyclass]
pub struct Node {
    events: Events,
    node: DoraNode,
}

#[pymethods]
impl Node {
    #[new]
    pub fn new() -> eyre::Result<Self> {
        let (node, events) = DoraNode::init_from_env()?;

        Ok(Node {
            events: Events::Dora(events),
            node,
        })
    }

    /// `.next()` gives you the next input that the node has received.
    /// It blocks until the next event becomes available.
    /// It will return `None` when all senders has been dropped.
    ///
    /// ```python
    /// event = node.next()
    /// ```
    ///
    /// You can also iterate over the event stream with a loop
    ///
    /// ```python
    /// for event in node:
    ///    match event["type"]:
    ///        case "INPUT":
    ///            match event["id"]:
    ///                 case "image":
    /// ```
    #[allow(clippy::should_implement_trait)]
    pub fn next(&mut self, py: Python) -> PyResult<Option<PyEvent>> {
        self.__next__(py)
    }

    pub fn __next__(&mut self, py: Python) -> PyResult<Option<PyEvent>> {
        let event = py.allow_threads(|| self.events.recv());
        Ok(event)
    }

    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    /// `send_output` send data from the node.
    ///
    /// ```python
    /// Args:
    ///    output_id: str,
    ///    data: Bytes|Arrow,
    ///    metadata: Option[Dict],
    /// ```
    ///
    /// ```python
    /// node.send_output("string", b"string", {"open_telemetry_context": "7632e76"})
    /// ```
    ///
    pub fn send_output(
        &mut self,
        output_id: String,
        data: PyObject,
        metadata: Option<&PyDict>,
        py: Python,
    ) -> eyre::Result<()> {
        let parameters = pydict_to_metadata(metadata)?;

        if let Ok(py_bytes) = data.downcast::<PyBytes>(py) {
            let data = py_bytes.as_bytes();
            self.node
                .send_output_bytes(output_id.into(), parameters, data.len(), data)
                .wrap_err("failed to send output")?;
        } else if let Ok(arrow_array) = arrow::array::ArrayData::from_pyarrow(data.as_ref(py)) {
            self.node.send_output(
                output_id.into(),
                parameters,
                arrow::array::make_array(arrow_array),
            )?;
        } else {
            eyre::bail!("invalid `data` type, must by `PyBytes` or arrow array")
        }

        Ok(())
    }

    /// Returns the full dataflow descriptor that this node is part of.
    ///
    /// This method returns the parsed dataflow YAML file.
    pub fn dataflow_descriptor(&self, py: Python) -> pythonize::Result<PyObject> {
        pythonize::pythonize(py, self.node.dataflow_descriptor())
    }

    pub fn merge_external_events(
        &mut self,
        subscription: &mut Ros2Subscription,
    ) -> eyre::Result<()> {
        let subscription = subscription.into_stream()?;
        let stream = futures::stream::poll_fn(move |cx| {
            let s = subscription.as_stream().map(|item| {
                match item.context("failed to read ROS2 message") {
                    Ok((value, _info)) => Python::with_gil(|py| {
                        value
                            .to_pyarrow(py)
                            .context("failed to convert value to pyarrow")
                            .unwrap_or_else(|err| PyErr::from(err).to_object(py))
                    }),
                    Err(err) => Python::with_gil(|py| PyErr::from(err).to_object(py)),
                }
            });
            futures::pin_mut!(s);
            s.poll_next_unpin(cx)
        });

        // take out the event stream and temporarily replace it with a dummy
        let events = std::mem::replace(
            &mut self.events,
            Events::Merged(Box::new(futures::stream::empty())),
        );
        // update self.events with the merged stream
        self.events = Events::Merged(events.merge_external_send(Box::pin(stream)));

        Ok(())
    }
}

enum Events {
    Dora(EventStream),
    Merged(Box<dyn Stream<Item = MergedEvent<PyObject>> + Unpin + Send>),
}

impl Events {
    fn recv(&mut self) -> Option<PyEvent> {
        match self {
            Events::Dora(events) => events.recv().map(PyEvent::from),
            Events::Merged(events) => futures::executor::block_on(events.next()).map(PyEvent::from),
        }
    }
}

impl<'a> MergeExternalSend<'a, PyObject> for Events {
    type Item = MergedEvent<PyObject>;

    fn merge_external_send(
        self,
        external_events: impl Stream<Item = PyObject> + Unpin + Send + 'a,
    ) -> Box<dyn Stream<Item = Self::Item> + Unpin + Send + 'a> {
        match self {
            Events::Dora(events) => events.merge_external_send(external_events),
            Events::Merged(events) => {
                let merged = events.merge_external_send(external_events);
                Box::new(merged.map(|event| match event {
                    MergedEvent::Dora(e) => MergedEvent::Dora(e),
                    MergedEvent::External(e) => MergedEvent::External(e.flatten()),
                }))
            }
        }
    }
}

impl Node {
    pub fn id(&self) -> String {
        self.node.id().to_string()
    }
}

/// Start a runtime for Operators
#[pyfunction]
pub fn start_runtime() -> eyre::Result<()> {
    dora_runtime::main().wrap_err("Dora Runtime raised an error.")
}

#[pymodule]
fn dora(py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(start_runtime, m)?)?;
    m.add_class::<Node>().unwrap();

    let ros2_bridge = PyModule::new(py, "ros2_bridge")?;
    dora_ros2_bridge_python::create_dora_ros2_bridge_module(ros2_bridge)?;
    let experimental = PyModule::new(py, "experimental")?;
    experimental.add_submodule(ros2_bridge)?;
    m.add_submodule(experimental)?;

    Ok(())
}