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, Event, EventStream};
use dora_operator_api_python::{metadata_to_pydict, pydict_to_metadata};
use eyre::Context;
use pyo3::exceptions::PyLookupError;
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict};

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

#[pyclass]
pub struct PyEvent(Event);

#[pymethods]
impl PyEvent {
    pub fn __getitem__(&mut self, key: &str, py: Python<'_>) -> PyResult<PyObject> {
        let value = match key {
            "type" => Some(self.ty().to_object(py)),
            "id" => self.id().map(|v| v.to_object(py)),
            "data" => self.data(py),
            "data_arrow" => self.data_arrow(py)?,
            "metadata" => self.metadata(py),
            "error" => self.error().map(|v| v.to_object(py)),
            other => {
                return Err(PyLookupError::new_err(format!(
                    "event has no property `{other}`"
                )))
            }
        };
        value.ok_or_else(|| PyLookupError::new_err(format!("event has no property `{key}`")))
    }
}

impl PyEvent {
    fn ty(&self) -> &str {
        match &self.0 {
            Event::Stop => "STOP",
            Event::Input { .. } => "INPUT",
            Event::InputClosed { .. } => "INPUT_CLOSED",
            Event::Error(_) => "ERROR",
            _other => "UNKNOWN",
        }
    }

    fn id(&self) -> Option<&str> {
        match &self.0 {
            Event::Input { id, .. } => Some(id),
            Event::InputClosed { id } => Some(id),
            _ => None,
        }
    }

    fn data(&self, py: Python<'_>) -> Option<PyObject> {
        match &self.0 {
            Event::Input {
                data: Some(data), ..
            } => Some(PyBytes::new(py, data).into()),
            _ => None,
        }
    }

    fn data_arrow(&mut self, py: Python<'_>) -> PyResult<Option<PyObject>> {
        if let Event::Input { data, .. } = &mut self.0 {
            if let Some(data) = data.take() {
                let array = data
                    .into_arrow_array()
                    .map_err(|err| arrow::pyarrow::PyArrowException::new_err(err.to_string()))?;
                // TODO: Does this call leak data?
                let array_data = array.to_pyarrow(py)?;
                return Ok(Some(array_data));
            }
        }
        Ok(None)
    }

    fn metadata(&self, py: Python<'_>) -> Option<PyObject> {
        match &self.0 {
            Event::Input { metadata, .. } => Some(metadata_to_pydict(metadata, py).to_object(py)),
            _ => None,
        }
    }

    fn error(&self) -> Option<&str> {
        match &self.0 {
            Event::Error(error) => Some(error),
            _other => None,
        }
    }
}

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