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

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

use arrow::pyarrow::PyArrowConvert;
use dora_node_api::{DoraNode, EventStream};
use dora_operator_api_python::{pydict_to_metadata, PyEvent};
use eyre::Context;
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict};

#[pyclass]
pub struct Node {
    events: EventStream,
    node: DoraNode,
}

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

        Ok(Node { events, node })
    }

    #[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.map(PyEvent))
    }

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

    pub fn send_output(
        &mut self,
        output_id: String,
        data: PyObject,
        metadata: Option<&PyDict>,
        py: Python,
    ) -> eyre::Result<()> {
        if let Ok(py_bytes) = data.downcast::<PyBytes>(py) {
            let data = py_bytes.as_bytes();
            self.send_output_slice(output_id, data.len(), data, metadata)
        } else if let Ok(arrow_array) = arrow::array::ArrayData::from_pyarrow(data.as_ref(py)) {
            if arrow_array.data_type() != &arrow::datatypes::DataType::UInt8 {
                eyre::bail!("only arrow arrays with data type `UInt8` are supported");
            }
            if arrow_array.buffers().len() != 1 {
                eyre::bail!("output arrow array must contain a single buffer");
            }

            let len = arrow_array.len();
            let slice = &arrow_array.buffer(0)[..len];

            self.send_output_slice(output_id, len, slice, metadata)
        } else {
            eyre::bail!("invalid `data` type, must by `PyBytes` or arrow array")
        }
    }
}

impl Node {
    fn send_output_slice(
        &mut self,
        output_id: String,
        len: usize,
        data: &[u8],
        metadata: Option<&PyDict>,
    ) -> eyre::Result<()> {
        let metadata = pydict_to_metadata(metadata)?;
        self.node
            .send_output(output_id.into(), metadata, len, |out| {
                out.copy_from_slice(data);
            })
            .wrap_err("failed to send output")
    }

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

#[pyfunction]
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();
    Ok(())
}