raphtory 0.17.0

//! Helper functions for the Python bindings.
//!
//! This module contains helper functions for the Python bindings.
//! These functions are not part of the public API and are not exported to the Python module.
use crate::{
    core::{
        entities::{
            nodes::node_ref::{AsNodeRef, NodeRef},
            GidRef,
        },
        storage::timeindex::AsTime,
    },
    db::api::view::*,
    python::graph::node::PyNode,
};
use chrono::{DateTime, Utc};
use numpy::{IntoPyArray, PyArray};
use pyo3::{exceptions::PyTypeError, prelude::*, pybacked::PyBackedStr, BoundObject};
use raphtory_api::core::entities::{
    properties::prop::{Prop, PropUnwrap},
    VID,
};
use std::{future::Future, sync::OnceLock};
use tokio::runtime::{Builder, Runtime};

pub mod errors;
pub(crate) mod export;
mod module_helpers;

#[derive(Debug, Eq, PartialEq, Hash)]
pub enum PyNodeRef {
    ExternalStr(PyBackedStr),
    ExternalInt(u64),
    Internal(VID),
}

impl<'source> FromPyObject<'source> for PyNodeRef {
    fn extract_bound(ob: &Bound<'source, PyAny>) -> PyResult<Self> {
        if let Ok(s) = ob.extract::<PyBackedStr>() {
            Ok(PyNodeRef::ExternalStr(s))
        } else if let Ok(gid) = ob.extract::<u64>() {
            Ok(PyNodeRef::ExternalInt(gid))
        } else if let Ok(v) = ob.extract::<PyNode>() {
            Ok(PyNodeRef::Internal(v.node.node))
        } else {
            Err(PyTypeError::new_err("Not a valid node"))
        }
    }
}

impl AsNodeRef for PyNodeRef {
    fn as_node_ref(&self) -> NodeRef<'_> {
        match self {
            PyNodeRef::ExternalStr(str) => NodeRef::External(GidRef::Str(str)),
            PyNodeRef::ExternalInt(gid) => NodeRef::External(GidRef::U64(*gid)),
            PyNodeRef::Internal(vid) => NodeRef::Internal(*vid),
        }
    }
}

// TODO: Revisit once the two lifetime version of FromPyObject is available in pyo3 (see https://github.com/PyO3/pyo3/pull/4390)
// /// Extract a `NodeRef` from a Python object.
// /// The object can be a `str`, `u64` or `PyNode`.
// /// If the object is a `PyNode`, the `NodeRef` is extracted from the `PyNode`.
// /// If the object is a `str`, the `NodeRef` is created from the `str`.
// /// If the object is a `int`, the `NodeRef` is created from the `int`.
// ///
// /// Arguments
// ///     vref: The Python object to extract the `NodeRef` from.
// ///
// /// Returns
// ///    A `NodeRef` extracted from the Python object.
// impl<'source> FromPyObject<'source> for NodeRef<'source> {
//     fn extract_bound(vref: &Bound<'source, PyAny>) -> PyResult<Self> {
//         if let Ok(s) = vref.extract::<&'source str>() {
//             Ok(NodeRef::External(GidRef::Str(s)))
//         } else if let Ok(gid) = vref.extract::<u64>() {
//             Ok(NodeRef::External(GidRef::U64(gid)))
//         } else if let Ok(v) = vref.extract::<PyNode>() {
//             Ok(NodeRef::Internal(v.node.node))
//         } else {
//             Err(PyTypeError::new_err("Not a valid node"))
//         }
//     }
// }

pub trait WindowSetOps {
    fn build_iter(&self) -> PyGenericIterator;
    fn time_index(&self, center: bool) -> PyGenericIterable;
}

impl<T> WindowSetOps for WindowSet<'static, T>
where
    T: TimeOps<'static> + Clone + Sync + Send + 'static,
    T::WindowedViewType: for<'py> IntoPyObject<'py> + Send + Sync + 'static,
{
    fn build_iter(&self) -> PyGenericIterator {
        self.clone().into()
    }

    fn time_index(&self, center: bool) -> PyGenericIterable {
        let window_set = self.clone();

        if window_set.temporal() {
            let iterable = move || {
                let iter: BoxedIter<DateTime<Utc>> = Box::new(
                    window_set
                        .clone()
                        .time_index(center)
                        .flat_map(|timestamp| timestamp.dt()),
                );
                iter
            };
            iterable.into()
        } else {
            (move || {
                let iter: BoxedIter<i64> = Box::new(window_set.time_index(center));
                iter
            })
            .into()
        }
    }
}

#[pyclass(name = "WindowSet", module = "raphtory", frozen)]
pub struct PyWindowSet {
    window_set: Box<dyn WindowSetOps + Send + Sync>,
}

impl<T> From<WindowSet<'static, T>> for PyWindowSet
where
    T: TimeOps<'static> + Clone + Sync + Send + 'static,
    T::WindowedViewType: for<'py> IntoPyObject<'py> + Send + Sync,
{
    fn from(value: WindowSet<'static, T>) -> Self {
        Self {
            window_set: Box::new(value),
        }
    }
}

impl<'py, T> IntoPyObject<'py> for WindowSet<'static, T>
where
    T: TimeOps<'static> + Clone + Sync + Send + 'static,
    T::WindowedViewType: for<'py2> IntoPyObject<'py2> + Send + Sync,
{
    type Target = PyWindowSet;
    type Output = <Self::Target as IntoPyObject<'py>>::Output;
    type Error = <Self::Target as IntoPyObject<'py>>::Error;

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        PyWindowSet::from(self).into_pyobject(py)
    }
}

#[pymethods]
impl PyWindowSet {
    fn __iter__(&self) -> PyGenericIterator {
        self.window_set.build_iter()
    }

    /// Returns the time index of this window set.
    ///
    /// It uses the last time of each window as the reference or the center of each if `center` is
    /// set to `True`.
    ///
    /// Arguments:
    ///     center (bool): If True time indexes are centered. Defaults to False.
    ///
    /// Returns:
    ///     Iterable: The time index.
    #[pyo3(signature = (center=false))]
    fn time_index(&self, center: bool) -> PyGenericIterable {
        self.window_set.time_index(center)
    }
}

#[pyclass(name = "Iterable")]
pub struct PyGenericIterable {
    build_iter: Box<dyn Fn() -> BoxedIter<PyResult<PyObject>> + Send + Sync>,
}

impl<F, I: Send + Sync, T> From<F> for PyGenericIterable
where
    F: (Fn() -> I) + Send + Sync + 'static,
    I: Iterator<Item = T> + Send + 'static,
    T: for<'py> IntoPyObject<'py> + 'static,
{
    fn from(value: F) -> Self {
        let build_py_iter: Box<dyn Fn() -> BoxedIter<PyResult<PyObject>> + Send + Sync> =
            Box::new(move || {
                Box::new(value().map(|item| {
                    Python::with_gil(|py| {
                        Ok(item
                            .into_pyobject(py)
                            .map_err(|e| e.into())?
                            .into_any()
                            .unbind())
                    })
                }))
            });
        Self {
            build_iter: build_py_iter,
        }
    }
}

#[pymethods]
impl PyGenericIterable {
    fn __iter__(&self) -> PyGenericIterator {
        PyGenericIterator::new((self.build_iter)())
    }
}

#[pyclass(name = "Iterator", unsendable)]
pub struct PyGenericIterator {
    iter: Box<dyn Iterator<Item = PyResult<PyObject>>>,
}

impl PyGenericIterator {
    pub fn new(iter: Box<dyn Iterator<Item = PyResult<PyObject>>>) -> Self {
        Self { iter }
    }
    pub fn from_result_iter<I, T, E>(iter: I) -> Self
    where
        I: Iterator<Item = Result<T, E>> + 'static,
        T: for<'py> IntoPyObject<'py> + 'static,
        PyErr: From<E>,
    {
        let py_iter = Box::new(iter.map(|result| {
            Python::with_gil(|py| match result {
                Ok(item) => Ok(item
                    .into_pyobject(py)
                    .map_err(|e| e.into())?
                    .into_any()
                    .unbind()),
                Err(time_error) => Err(PyErr::from(time_error)),
            })
        }));
        Self { iter: py_iter }
    }
}

impl<I, T> From<I> for PyGenericIterator
where
    I: Iterator<Item = T> + 'static,
    T: for<'py> IntoPyObject<'py> + 'static,
{
    fn from(value: I) -> Self {
        let py_iter = Box::new(value.map(|item| {
            Python::with_gil(|py| {
                Ok(item
                    .into_pyobject(py)
                    .map_err(|e| e.into())?
                    .into_any()
                    .unbind())
            })
        }));
        Self { iter: py_iter }
    }
}

impl IntoIterator for PyGenericIterator {
    type Item = PyResult<PyObject>;

    type IntoIter = Box<dyn Iterator<Item = Self::Item>>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter
    }
}

#[pymethods]
impl PyGenericIterator {
    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }
    fn __next__(&mut self) -> Option<PyResult<PyObject>> {
        self.iter.next()
    }
}

#[pyclass(name = "NestedIterator")]
pub struct PyNestedGenericIterator {
    iter: BoxedIter<PyGenericIterator>,
}

impl PyNestedGenericIterator {
    pub fn from_nested_result_iter<I, J, T, E>(iter: I) -> Self
    where
        I: Iterator<Item = J> + Send + Sync + 'static,
        J: Iterator<Item = Result<T, E>> + Send + Sync + 'static,
        T: for<'py> IntoPyObject<'py> + 'static,
        PyErr: From<E>,
    {
        let py_iter = Box::new(iter.map(|item| PyGenericIterator::from_result_iter(item)));
        Self { iter: py_iter }
    }
}

impl<I, J, T> From<I> for PyNestedGenericIterator
where
    I: Iterator<Item = J> + Send + Sync + 'static,
    J: Iterator<Item = T> + Send + Sync + 'static,
    T: for<'py> IntoPyObject<'py> + 'static,
{
    fn from(value: I) -> Self {
        let py_iter = Box::new(value.map(|item| item.into()));
        Self { iter: py_iter }
    }
}

#[pymethods]
impl PyNestedGenericIterator {
    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }
    fn __next__(&mut self) -> Option<PyGenericIterator> {
        self.iter.next()
    }
}

pub enum NumpyArray {
    Bool(Vec<bool>),
    U32(Vec<u32>),
    U64(Vec<u64>),
    I32(Vec<i32>),
    I64(Vec<i64>),
    F32(Vec<f32>),
    F64(Vec<f64>),
    Props(Vec<Prop>),
}

impl FromIterator<Prop> for NumpyArray {
    fn from_iter<I: IntoIterator<Item = Prop>>(iter: I) -> Self {
        let mut iter = iter.into_iter().peekable();
        match iter.peek() {
            Some(Prop::Bool(_)) => Self::Bool(iter.filter_map(|p| p.into_bool()).collect()),
            Some(Prop::I32(_)) => Self::I32(iter.filter_map(|p| p.into_i32()).collect()),
            Some(Prop::I64(_)) => Self::I64(iter.filter_map(|p| p.into_i64()).collect()),
            Some(Prop::U32(_)) => Self::U32(iter.filter_map(|p| p.into_u32()).collect()),
            Some(Prop::U64(_)) => Self::U64(iter.filter_map(|p| p.into_u64()).collect()),
            Some(Prop::F32(_)) => Self::F32(iter.filter_map(|p| p.into_f32()).collect()),
            Some(Prop::F64(_)) => Self::F64(iter.filter_map(|p| p.into_f64()).collect()),
            _ => Self::Props(iter.collect()),
        }
    }
}

impl From<Vec<i64>> for NumpyArray {
    fn from(value: Vec<i64>) -> Self {
        NumpyArray::I64(value)
    }
}

impl<'py> IntoPyObject<'py> for NumpyArray {
    type Target = PyAny;
    type Output = Bound<'py, Self::Target>;
    type Error = PyErr;

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        match self {
            NumpyArray::Bool(value) => Ok(value.into_pyarray(py).into_any()),
            NumpyArray::I32(value) => Ok(value.into_pyarray(py).into_any()),
            NumpyArray::I64(value) => Ok(value.into_pyarray(py).into_any()),
            NumpyArray::U32(value) => Ok(value.into_pyarray(py).into_any()),
            NumpyArray::U64(value) => Ok(value.into_pyarray(py).into_any()),
            NumpyArray::F32(value) => Ok(value.into_pyarray(py).into_any()),
            NumpyArray::F64(value) => Ok(value.into_pyarray(py).into_any()),
            NumpyArray::Props(vec) => match vec.first() {
                Some(Prop::Bool(_)) => Ok(PyArray::from_iter(
                    py,
                    vec.into_iter().filter_map(|p| p.into_bool()),
                )
                .into_any()),
                Some(Prop::I32(_)) => Ok(PyArray::from_iter(
                    py,
                    vec.into_iter().filter_map(|p| p.into_i32()),
                )
                .into_any()),
                Some(Prop::I64(_)) => Ok(PyArray::from_iter(
                    py,
                    vec.into_iter().filter_map(|p| p.into_i64()),
                )
                .into_any()),
                Some(Prop::U32(_)) => Ok(PyArray::from_iter(
                    py,
                    vec.into_iter().filter_map(|p| p.into_u32()),
                )
                .into_any()),
                Some(Prop::U64(_)) => Ok(PyArray::from_iter(
                    py,
                    vec.into_iter().filter_map(|p| p.into_u64()),
                )
                .into_any()),
                Some(Prop::F32(_)) => Ok(PyArray::from_iter(
                    py,
                    vec.into_iter().filter_map(|p| p.into_f32()),
                )
                .into_any()),
                Some(Prop::F64(_)) => Ok(PyArray::from_iter(
                    py,
                    vec.into_iter().filter_map(|p| p.into_f64()),
                )
                .into_any()),
                _ => vec.into_pyobject(py),
            },
        }
    }
}

// This function takes a function that returns a future instead of taking just a future because
// a task might return an unsendable future but what we can do is making a function returning that
// future which is sendable itself
pub(crate) fn execute_async_task<T, F, O>(task: T) -> O
where
    T: FnOnce() -> F + Send + 'static,
    F: Future<Output = O> + 'static,
    O: Send + 'static,
{
    Python::with_gil(|py| py.allow_threads(move || get_runtime().block_on(task())))
}

static RUNTIME: OnceLock<Runtime> = OnceLock::new();

pub fn get_runtime() -> &'static Runtime {
    RUNTIME.get_or_init(|| {
        Builder::new_multi_thread()
            .enable_all()
            // Optional: limit threads if you want to leave room for Python
            .worker_threads(4)
            .build()
            .expect("Failed to create Tokio runtime")
    })
}

pub fn block_on<F: Future>(future: F) -> F::Output {
    get_runtime().block_on(future)
}