dataloader-rs 0.3.0

use std::slice;

use crate::{dataset::Dataset, error::Result, python::pool::attach};
use pyo3::exceptions::{PyNotImplementedError, PyTypeError, PyValueError};
use pyo3::intern;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList, PyTuple};

#[pyclass(frozen, subclass)]
pub struct PyDatasetBase;

#[pymethods]
impl PyDatasetBase {
    #[new]
    #[pyo3(signature = (*_args, **_kwargs))]
    fn new(_args: &Bound<'_, PyTuple>, _kwargs: Option<&Bound<'_, PyDict>>) -> Self {
        Self
    }

    fn __len__(&self) -> PyResult<usize> {
        Err(PyNotImplementedError::new_err(
            "PyDataset.__len__ must be implemented by subclasses",
        ))
    }

    fn __getitem__(&self, _index: usize) -> PyResult<Py<PyAny>> {
        Err(PyNotImplementedError::new_err(
            "PyDataset.__getitem__ must be implemented by subclasses",
        ))
    }
}

pub(crate) type PyDataset = Py<PyDatasetBase>;

pub(crate) fn len_py(dataset: &PyDataset, py: Python<'_>) -> PyResult<usize> {
    dataset
        .call_method0(py, intern!(py, "__len__"))?
        .extract(py)
}

pub(crate) fn get_item_py(
    dataset: &PyDataset,
    py: Python<'_>,
    index: usize,
) -> PyResult<Py<PyAny>> {
    let item = dataset.call_method1(py, intern!(py, "__getitem__"), (index,))?;
    check_not_awaitable(py, slice::from_ref(&item))?;
    Ok(item)
}

/// Refuse awaitable items: a sync `__getitem__` that returns a coroutine (an
/// `async def` helper called without `await`) would fill batches with
/// coroutines that never run. Checked on the first item of `items`; the
/// coroutines are then closed so that Python does not warn that they were
/// never awaited.
fn check_not_awaitable(py: Python<'_>, items: &[Py<PyAny>]) -> PyResult<()> {
    let Some(first) = items.first() else {
        return Ok(());
    };
    let awaitable = intern!(py, "__await__");
    if first.bind(py).get_type().getattr_opt(awaitable)?.is_none() {
        return Ok(());
    }
    for item in items {
        if let Some(close) = item.bind(py).getattr_opt(intern!(py, "close"))? {
            close.call0()?;
        }
    }
    Err(PyTypeError::new_err(
        "PyDataset.__getitem__ returned an awaitable; subclass PyAsyncDataset \
         and define `async def __getitem__` instead",
    ))
}

/// How a batch of samples is fetched from a Python dataset.
///
/// Mirrors PyTorch's fetcher: a dataset defining `__getitems__(indices)` is
/// called once per batch with the list of indices and returns the list of
/// samples; otherwise `__getitem__` is called once per index. As in PyTorch,
/// a falsy `__getitems__` (e.g. set to `None`) counts as not defined.
pub(crate) enum Fetcher {
    /// Bound `__getitems__` method.
    Batched(Py<PyAny>),
    /// Bound `__getitem__` method, cached to skip one lookup per item.
    PerItem(Py<PyAny>),
}

impl Fetcher {
    pub(crate) fn new(dataset: &PyDataset, py: Python<'_>) -> PyResult<Self> {
        let dataset = dataset.bind(py);
        Ok(match dataset.getattr_opt(intern!(py, "__getitems__"))? {
            Some(getitems) if getitems.is_truthy()? => Self::Batched(getitems.unbind()),
            _ => Self::PerItem(dataset.getattr(intern!(py, "__getitem__"))?.unbind()),
        })
    }

    /// The bound method this fetcher calls.
    pub(crate) fn method(&self) -> &Py<PyAny> {
        match self {
            Self::Batched(method) | Self::PerItem(method) => method,
        }
    }

    pub(crate) fn fetch(&self, py: Python<'_>, indices: &[usize]) -> PyResult<Vec<Py<PyAny>>> {
        match self {
            Self::PerItem(getitem) => {
                let items = indices
                    .iter()
                    .map(|&i| getitem.call1(py, (i,)))
                    .collect::<PyResult<Vec<_>>>()?;
                check_not_awaitable(py, &items)?;
                Ok(items)
            }
            Self::Batched(getitems) => {
                let samples = getitems.bind(py).call1((PyList::new(py, indices)?,))?;
                let samples = samples
                    .try_iter()?
                    .map(|sample| sample.map(Bound::unbind))
                    .collect::<PyResult<Vec<_>>>()?;
                // A short or long list would silently change batch sizes.
                if samples.len() != indices.len() {
                    return Err(PyValueError::new_err(format!(
                        "__getitems__ returned {} samples for {} indices",
                        samples.len(),
                        indices.len()
                    )));
                }
                check_not_awaitable(py, &samples)?;
                Ok(samples)
            }
        }
    }
}

impl Dataset for PyDataset {
    type Item = Py<PyAny>;

    fn get(&self, index: usize) -> Result<Self::Item> {
        Ok(attach(|py| get_item_py(self, py, index))?)
    }

    /// Acquire the Python thread state once for the entire batch rather than
    /// once per item — reduces GIL attach/release overhead from O(batch_size)
    /// to O(1) per batch in the threaded (num_workers>0) code path.
    ///
    /// Uses `__getitems__` when the dataset defines it, else a cached bound
    /// `__getitem__` (one attribute lookup per batch, not per item).
    fn get_batch(&self, indices: &[usize]) -> Result<Vec<Py<PyAny>>> {
        Ok(attach(|py| Fetcher::new(self, py)?.fetch(py, indices))?)
    }

    fn len(&self) -> usize {
        Python::attach(|py| {
            len_py(self, py).unwrap_or_else(|e| panic!("PyDataset.__len__ failed: {e}"))
        })
    }
}