dataloader-rs 0.2.0

use std::time::{Duration, Instant};

use pyo3::prelude::*;
use pyo3::{PyTraverseError, PyVisit};

use crate::python::collator::PyCollator;
use crate::python::dataloader::PyDataloader;
use crate::python::dataset::{Fetcher, PyDataset};
use crate::python::into_py_err;

type CorePyDataloaderIter = crate::loader::DataLoaderIter<PyDataset, PyCollator>;

/// How long `__next__` waits with the interpreter released before it runs
/// signal handlers: Ctrl-C ends a wait within this.
const SIGNAL_CHECK_INTERVAL: Duration = Duration::from_millis(100);

pub(crate) enum PyIterInner {
    /// Sequential (num_workers=0): call Python directly in `__next__` using the
    /// already-held `py` token — no thread, no channel, no extra GIL acquire.
    Direct {
        chunks: std::vec::IntoIter<Vec<usize>>,
        /// `__getitems__` or cached `__getitem__`, resolved once per epoch.
        fetcher: Fetcher,
        /// Collator cloned at `__iter__` time so `__next__` skips borrowing `_owner`.
        collator: PyCollator,
    },
    /// Parallel (num_workers>0): batch tasks on the loader's pool.
    Threaded(CorePyDataloaderIter),
}

#[pyclass(name = "PyDataloaderIter", module = "dataloader_rs")]
pub struct PyDataloaderIter {
    /// Keeps the `PyDataloader` alive while the iterator exists.
    pub(crate) _owner: Py<PyDataloader>,
    pub(crate) inner: PyIterInner,
}

#[pymethods]
impl PyDataloaderIter {
    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    fn __next__(&mut self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
        let batch = match &mut self.inner {
            PyIterInner::Direct {
                chunks,
                fetcher,
                collator,
            } => {
                let Some(chunk) = chunks.next() else {
                    return Ok(None);
                };
                // We already hold `py` — call Python directly with no extra GIL acquire.
                let items = fetcher.fetch(py, &chunk)?;
                collator.collate_with_py(py, items).map_err(into_py_err)?
            }

            PyIterInner::Threaded(inner) => {
                // A batch that already arrived is taken without releasing the
                // interpreter. Otherwise wait with it released, so workers can
                // attach, in slices: between two, signal handlers run, so
                // Ctrl-C raises `KeyboardInterrupt` here and the iterator stays
                // usable. A worker's panic unwinds from here to PyO3, which
                // raises it; the epoch is then over.
                let next = match inner.next_until(None) {
                    Some(next) => next,
                    None => loop {
                        let deadline = Instant::now() + SIGNAL_CHECK_INTERVAL;
                        if let Some(next) = py.detach(|| inner.next_until(Some(deadline))) {
                            break next;
                        }
                        py.check_signals()?;
                    },
                };
                let Some(batch) = next else {
                    return Ok(None);
                };
                batch.map_err(into_py_err)?
            }
        };
        batch.into_py(py).map(Some)
    }

    /// The loader, plus the references the direct path holds itself; a
    /// threaded epoch shares the loader's objects.
    fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
        visit.call(&self._owner)?;
        if let PyIterInner::Direct {
            fetcher, collator, ..
        } = &self.inner
        {
            visit.call(fetcher.method())?;
            visit.call(collator.collate_fn())?;
        }
        Ok(())
    }

    fn __len__(&self) -> usize {
        match &self.inner {
            PyIterInner::Direct { chunks, .. } => chunks.len(),
            PyIterInner::Threaded(inner) => inner.len(),
        }
    }
}