kaparoo-python 0.6.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

kaparoo/data/sequences/base.py

@@ -1,3 +1,5 @@
+"""The `DataSequence[T, M]` abstract base: indexable items with metadata."""
+
 from __future__ import annotations
 
 __all__ = ("DataSequence",)
@@ -32,7 +34,7 @@ class DataSequence[T, M = None](Sequence[T]):
 
     @abstractmethod
     def __len__(self) -> int:
-        raise NotImplementedError
+        """Return the number of items in the sequence."""
 
     # --- item access -------------------------------------------------------
 
@@ -50,24 +52,36 @@ class DataSequence[T, M = None](Sequence[T]):
 
     @abstractmethod
     def get_item(self, index: int) -> T:
-        raise NotImplementedError
+        """Fetch and return the item at `index`."""
 
     def get_items(self, indices: Sequence[int]) -> Sequence[T]:
+        """Fetch many items at once, in `indices` order.
+
+        Defaults to one `get_item` per index; override to use a backing
+        store's native batch read.
+        """
         return [self.get_item(index) for index in indices]
 
     # --- metadata access ---------------------------------------------------
 
     @abstractmethod
     def get_meta(self, index: int) -> M:
-        raise NotImplementedError
+        """Return the metadata for the item at `index` (`None` when `M` is `None`)."""
 
     def get_metas(self, indices: Sequence[int]) -> Sequence[M]:
+        """Fetch many metadata values at once, in `indices` order.
+
+        Defaults to one `get_meta` per index; override alongside
+        `get_items` when a batch read is cheaper.
+        """
         return [self.get_meta(index) for index in indices]
 
     # --- combined item + metadata ------------------------------------------
 
     def get_pair(self, index: int) -> tuple[T, M]:
+        """Return the `(item, metadata)` pair at `index`."""
         return self.get_item(index), self.get_meta(index)
 
     def get_pairs(self, indices: Sequence[int]) -> Sequence[tuple[T, M]]:
+        """Fetch many `(item, metadata)` pairs at once, in `indices` order."""
         return [self.get_pair(index) for index in indices]

kaparoo/data/sequences/composers.py

@@ -1,3 +1,5 @@
+"""Lazy `DataSequence` composers: slice, concat, transform, window, zip."""
+
 from __future__ import annotations
 
 __all__ = (
@@ -10,7 +12,7 @@ __all__ = (
 
 from abc import abstractmethod
 from bisect import bisect_right
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, cast, override
 
 from kaparoo.data.sequences.base import DataSequence
 
@@ -18,23 +20,35 @@ if TYPE_CHECKING:
     from collections.abc import Callable, Sequence
 
 
-class SlicedSequence[T, M](DataSequence[T, M]):
-    """A view of `source` exposing only items at the given `indices`.
+def _resolve_index(index: int, length: int) -> int:
+    """Normalize a possibly-negative index against `length`, validating range.
+
+    Used by `ConcatSequence`, `WindowedSequence`, and `ZippedSequence`.
+    `SlicedSequence` intentionally opts out -- it indexes its `indices` tuple
+    directly, which wraps and raises the same way but with the builtin message.
+
+    Raises:
+        IndexError: If `index` is outside `[-length, length)`.
+    """
+    original = index
+    if index < 0:
+        index += length
+    if not 0 <= index < length:
+        msg = f"index {original} out of range for length {length}"
+        raise IndexError(msg)
+    return index
+
 
-    `indices` is materialized as a tuple at construction time so that the
-    view has a stable length and supports O(1) random access. Negative
-    and out-of-range indices delegate to Python's tuple semantics
-    (negative wraps, out-of-range raises `IndexError`).
+class SlicedSequence[T, M](DataSequence[T, M]):
+    """A view of `source` exposing only the items at `indices`, in that order.
 
-    `indices` is taken as-is: duplicates are allowed (the same source
-    item is yielded multiple times) and order is preserved (no sorting).
-    Bounds against `source` are not validated at construction; an
-    out-of-range entry surfaces only when that position is accessed.
+    `indices` is taken as-is -- duplicates repeat the source item, order is
+    preserved -- and is not bounds-checked against `source` until a position
+    is accessed. A negative view index wraps; out of range raises `IndexError`.
 
     Example:
         >>> sliced = SlicedSequence(full_dataset, [3, 7, 11])
         >>> sliced[0]  # == full_dataset[3]
-        >>> sliced[1]  # == full_dataset[7]
     """
 
     def __init__(
@@ -58,42 +72,54 @@ class SlicedSequence[T, M](DataSequence[T, M]):
     def __len__(self) -> int:
         return len(self._indices)
 
+    @override
     def get_item(self, index: int) -> T:
+        """Fetch the source item at the mapped index `indices[index]`."""
         return self._source.get_item(self._indices[index])
 
+    @override
     def get_meta(self, index: int) -> M:
+        """Fetch the source metadata at the mapped index `indices[index]`."""
         return self._source.get_meta(self._indices[index])
 
+    @override
+    def get_items(self, indices: Sequence[int]) -> Sequence[T]:
+        """Map each view index through `indices`, then batch-fetch from `source`."""
+        return self._source.get_items([self._indices[i] for i in indices])
+
+    @override
+    def get_metas(self, indices: Sequence[int]) -> Sequence[M]:
+        """Map each view index through `indices`, then batch-fetch metadata."""
+        return self._source.get_metas([self._indices[i] for i in indices])
+
 
 class TransformedSequence[T_in, M_in, T_out = T_in, M_out = M_in](
     DataSequence[T_out, M_out]
 ):
     """A view of `source` with `transform` applied lazily to each item.
 
-    `transform` is called on demand in `get_item`; nothing is loaded or
-    converted at construction time. `get_meta` passes through
-    `source.get_meta` unchanged by default -- override it in a subclass
-    when `M_out` differs from `M_in`.
+    `transform` runs on demand in `get_item`; nothing is converted at
+    construction. `get_meta` passes `source.get_meta` through unchanged, which
+    is correct only when the metadata type is unchanged (the default
+    `M_out == M_in`). **Override `get_meta` whenever `M_out != M_in`**: the
+    passthrough's `cast` cannot catch a missing override -- generics are erased
+    at runtime -- so a forgotten one silently yields an `M_in` value mistyped
+    as `M_out`.
 
     Type Parameters:
-        T_in: Item type of `source`.
-        M_in: Metadata type of `source`.
-        T_out: Item type after the transform. Defaults to `T_in`.
-        M_out: Metadata type exposed by this view. Defaults to `M_in`.
-            When `M_out != M_in`, override `get_meta` in a subclass;
-            the default passthrough is only safe when `M_out == M_in`.
+        T_in, M_in: The source's element and metadata types.
+        T_out: The transformed element type. Defaults to `T_in`.
+        M_out: The transformed metadata type. Defaults to `M_in` (the
+            passthrough case); set it and override `get_meta` otherwise.
 
     Example:
         >>> # Item-only transform; metadata passes through unchanged.
         >>> normalized = TransformedSequence(image_folder, normalize)
 
-        >>> # Meta transform via subclassing:
+        >>> # Metadata transform via subclassing:
         >>> class Augmented(TransformedSequence[ndarray, Path, ndarray, AugMeta]):
         ...     def get_meta(self, index: int) -> AugMeta:
-        ...         return AugMeta(
-        ...             path=self.source.get_meta(index),
-        ...             applied="normalize",
-        ...         )
+        ...         return AugMeta(self.source.get_meta(index), applied="normalize")
     """
 
     def __init__(
@@ -112,21 +138,33 @@ class TransformedSequence[T_in, M_in, T_out = T_in, M_out = M_in](
     def __len__(self) -> int:
         return len(self._source)
 
+    @override
     def get_item(self, index: int) -> T_out:
+        """Fetch the source item at `index` and apply `transform`."""
         return self._transform(self._source.get_item(index))
 
+    @override
+    def get_items(self, indices: Sequence[int]) -> Sequence[T_out]:
+        """Batch-fetch from `source` and apply `transform` to each item."""
+        return [self._transform(item) for item in self._source.get_items(indices)]
+
+    @override
     def get_meta(self, index: int) -> M_out:
-        # Passthrough by default. Override when M_out != M_in.
+        """Pass `source`'s metadata through unchanged (valid only when `M_out == M_in`)."""
+        # Passthrough -- correct only when M_out == M_in. A subclass with a
+        # different M_out MUST override this; the cast cannot catch a missing
+        # override, since generics are erased at runtime.
         return cast("M_out", self._source.get_meta(index))
 
 
 class ConcatSequence[T, M](DataSequence[T, M]):
     """The end-to-end concatenation of zero or more `sources`.
 
-    Indexing maps to `(source, local_index)` via a precomputed cumulative
-    length array and `bisect`, so a lookup is O(log N) in the number of
-    sources. Negative indices are normalized; out-of-range indices raise
-    `IndexError`.
+    A logical index maps to the `(source, local index)` it falls in -- an
+    O(log N) lookup in the number of sources. Negative indices wrap; out of
+    range raises `IndexError`. Batch access (`get_items` / `get_metas`)
+    delegates one grouped call per source, so a source's own batch
+    optimization is used, with results kept in request order.
 
     Example:
         >>> combined = ConcatSequence(train_a, train_b, train_c)
@@ -148,57 +186,89 @@ class ConcatSequence[T, M](DataSequence[T, M]):
     def __len__(self) -> int:
         return self._cumulative[-1]
 
-    def _locate(self, index: int) -> tuple[DataSequence[T, M], int]:
-        """Resolve a logical index to `(source, local_index)`.
+    def _locate_index(self, index: int) -> tuple[int, int]:
+        """Resolve a logical index to `(source position, local index)`.
 
         Raises:
             IndexError: If `index` is outside `[-len(self), len(self))`.
         """
-        n = self._cumulative[-1]
-        original = index
-        if index < 0:
-            index += n
-        if not 0 <= index < n:
-            msg = f"index {original} out of range for length {n}"
-            raise IndexError(msg)
+        index = _resolve_index(index, self._cumulative[-1])
         i = bisect_right(self._cumulative, index) - 1
-        return self._sources[i], index - self._cumulative[i]
+        return i, index - self._cumulative[i]
+
+    def _locate(self, index: int) -> tuple[DataSequence[T, M], int]:
+        """Resolve a logical index to `(source, local_index)`."""
+        i, local = self._locate_index(index)
+        return self._sources[i], local
 
+    def _gather[R](
+        self,
+        indices: Sequence[int],
+        fetch: Callable[[DataSequence[T, M], list[int]], Sequence[R]],
+    ) -> list[R]:
+        """Batch-fetch `indices` with one grouped `fetch` per source.
+
+        The shared core of `get_items` / `get_metas`, which differ only in the
+        per-source `fetch`; results are scattered back into request order.
+        """
+        buckets: dict[int, list[tuple[int, int]]] = {}
+        for position, index in enumerate(indices):
+            source_index, local = self._locate_index(index)
+            buckets.setdefault(source_index, []).append((position, local))
+
+        gathered: dict[int, R] = {}
+        for source_index, entries in buckets.items():
+            fetched = fetch(
+                self._sources[source_index], [local for _, local in entries]
+            )
+            for (position, _), value in zip(entries, fetched, strict=True):
+                gathered[position] = value
+        return [gathered[position] for position in range(len(indices))]
+
+    @override
     def get_item(self, index: int) -> T:
+        """Locate the source for `index` and fetch its local item."""
         source, local = self._locate(index)
         return source.get_item(local)
 
+    @override
+    def get_items(self, indices: Sequence[int]) -> Sequence[T]:
+        """Group `indices` by source and batch-fetch items, kept in request order."""
+        return self._gather(indices, lambda source, locals_: source.get_items(locals_))
+
+    @override
     def get_meta(self, index: int) -> M:
+        """Locate the source for `index` and fetch its local metadata."""
         source, local = self._locate(index)
         return source.get_meta(local)
 
+    @override
+    def get_metas(self, indices: Sequence[int]) -> Sequence[M]:
+        """Group `indices` by source and batch-fetch metadata, kept in request order."""
+        return self._gather(indices, lambda source, locals_: source.get_metas(locals_))
+
 
 class WindowedSequence[T, M_in, M_out = M_in](DataSequence[tuple[T, ...], M_out]):
     """An abstract sliding-window view over `source`.
 
-    Each item is a tuple of `size` items from `source`, starting at
-    position `i * step`, with intra-window stride `skip`. Indexed item
-    access (`get_item`) is implemented; **the window's metadata
-    strategy is intentionally left abstract** so the relationship
-    between per-frame `M_in` and window-level `M_out` is decided at
-    subclass-definition time.
-
-    Subclasses use the `source`, `size`, `step`, `skip` properties and
-    should call `_normalize_index` from `get_meta` so negative and
-    out-of-range window indices behave the same way as in `get_item`.
+    Each item is a tuple of `size` items from `source`, the window starting at
+    `i * step` with intra-window stride `skip`. `get_item` is implemented;
+    **the window's metadata is intentionally left abstract** so a subclass
+    decides how per-frame metadata becomes window metadata (`M_in` -> `M_out`).
+    Subclasses should call `_normalize_index` in their `get_meta` so window
+    indices behave as in `get_item`.
 
     Type Parameters:
-        T: Item type of `source` (also the per-frame type within each
-            window).
-        M_in: Metadata type of `source` (per-frame metadata).
-        M_out: Metadata type of the window. Defaults to `M_in`.
-            Determined by the subclass's `get_meta` return.
+        T: The source's element type; each item is a `tuple[T, ...]`.
+        M_in: The source's per-frame metadata type.
+        M_out: The window's metadata type a subclass produces. Defaults to
+            `M_in`.
 
     Args:
         source: The sequence to window over.
-        size: Number of items per window. Must be positive.
-        step: Position advance between consecutive windows. Defaults
-            to 1 (overlapping windows by `size - 1`).
+        size: Items per window. Must be positive.
+        step: Advance between consecutive windows. Defaults to 1 (windows
+            overlap by `size - 1`).
         skip: Intra-window stride. Defaults to 1 (consecutive frames).
 
     Raises:
@@ -261,25 +331,19 @@ class WindowedSequence[T, M_in, M_out = M_in](DataSequence[tuple[T, ...], M_out]
         Raises:
             IndexError: If `index` is outside `[-len(self), len(self))`.
         """
-        n = self._length
-        original = index
-        if index < 0:
-            index += n
-        if not 0 <= index < n:
-            msg = f"index {original} out of range for length {n}"
-            raise IndexError(msg)
-        return index
+        return _resolve_index(index, self._length)
 
+    @override
     def get_item(self, index: int) -> tuple[T, ...]:
+        """Build the window at `index` as a tuple of `size` strided source items."""
         index = self._normalize_index(index)
         start = index * self._step
-        return tuple(
-            self._source.get_item(start + j * self._skip) for j in range(self._size)
-        )
+        stop = start + self._size * self._skip
+        return tuple(self._source.get_items(range(start, stop, self._skip)))
 
     @abstractmethod
     def get_meta(self, index: int) -> M_out:
-        raise NotImplementedError
+        """Return the metadata for window `index` (the `M_in` -> `M_out` policy)."""
 
 
 class ZippedSequence[T1, T2, M1 = None, M2 = None](
@@ -291,26 +355,16 @@ class ZippedSequence[T1, T2, M1 = None, M2 = None](
     `(first.get_meta(i), second.get_meta(i))` -- the "paired image + label"
     pattern that `ConcatSequence` (end-to-end) cannot express.
 
-    With `strict=True` (the default) the two sequences must have the same
-    length; a mismatch raises `ValueError` at construction. With
-    `strict=False` the view is truncated to the shorter length, like the
-    builtin `zip`. For a different combined-metadata shape, subclass and
-    override `get_meta`.
+    With `strict=True` (the default) the sequences must be the same length, or
+    construction raises `ValueError`; with `strict=False` the view truncates to
+    the shorter, like the builtin `zip`. For a different combined-metadata
+    shape, subclass and override `get_meta`.
 
     Type Parameters:
-        T1: Item type of the first source.
-        T2: Item type of the second source.
-        M1: Metadata type of the first source. Defaults to `None`.
-        M2: Metadata type of the second source. Defaults to `None`.
-
-    Args:
-        first: The first sequence.
-        second: The second sequence.
-        strict: When True (default), require equal lengths and raise on a
-            mismatch. When False, truncate to the shorter length.
-
-    Raises:
-        ValueError: If `strict` is True and the sequences differ in length.
+        T1, T2: Element types of the first and second sequence; items are
+            `tuple[T1, T2]`.
+        M1, M2: Their metadata types; metadata is `tuple[M1, M2]`. Each
+            defaults to `None` (a sequence without metadata).
 
     Example:
         >>> pairs = ZippedSequence(images, labels)
@@ -354,20 +408,17 @@ class ZippedSequence[T1, T2, M1 = None, M2 = None](
         Raises:
             IndexError: If `index` is outside `[-len(self), len(self))`.
         """
-        n = self._length
-        original = index
-        if index < 0:
-            index += n
-        if not 0 <= index < n:
-            msg = f"index {original} out of range for length {n}"
-            raise IndexError(msg)
-        return index
+        return _resolve_index(index, self._length)
 
+    @override
     def get_item(self, index: int) -> tuple[T1, T2]:
+        """Fetch the paired `(first[index], second[index])` item."""
         index = self._normalize_index(index)
         return self._first.get_item(index), self._second.get_item(index)
 
+    @override
     def get_items(self, indices: Sequence[int]) -> Sequence[tuple[T1, T2]]:
+        """Normalize indices, then batch-fetch and pair items from both sources."""
         # Normalize, then bulk-delegate so each source's `get_items`
         # optimization is used.
         normalized = [self._normalize_index(i) for i in indices]
@@ -379,11 +430,15 @@ class ZippedSequence[T1, T2, M1 = None, M2 = None](
             )
         )
 
+    @override
     def get_meta(self, index: int) -> tuple[M1, M2]:
+        """Fetch the paired `(first, second)` metadata at `index`."""
         index = self._normalize_index(index)
         return self._first.get_meta(index), self._second.get_meta(index)
 
+    @override
     def get_metas(self, indices: Sequence[int]) -> Sequence[tuple[M1, M2]]:
+        """Normalize indices, then batch-fetch and pair metadata from both sources."""
         normalized = [self._normalize_index(i) for i in indices]
         return list(
             zip(