feat: strict typing stream.iter_array

docs/releases/unreleased.md

@@ -7,3 +7,7 @@
 - `stream.iter_csv` now closes the file it opened and restores `csv.field_size_limit` even when the stream is not exhausted, e.g. when the caller breaks out of the loop. Only the file `iter_csv` opened itself is closed; a buffer passed in by the caller is still left open.
 - `stream.iter_sql` now closes the result it iterates, so the underlying cursor is released once the stream is exhausted or abandoned.
 - `stream.cache`, `stream.iter_csv`, and `stream.iter_sql` are now clean under strict mypy. `sqlalchemy` is type-checked rather than ignored, so the `query` and `conn` arguments of `stream.iter_sql` are checked against the SQLAlchemy 2.0 types.
+- Enabled strict typing for `stream.iter_array`. It is now overloaded on its inputs: an `X` of texts is typed to yield `str` rows, and a plain sequence `y` propagates its own target type, while numpy arrays keep yielding `Any`. Typing-only, apart from the fixes below.
+- Fixed `stream.iter_array` for plain Python inputs, which are documented as supported: `shuffle=True` no longer raises a `TypeError` on lists, and a multi-output `y` given as a list of lists no longer raises an `AttributeError`.
+- `stream.iter_array` now yields an empty stream when `X` is empty, instead of raising an `IndexError`, and comes out slightly faster on every input shape.
+- Added dedicated tests for `stream.iter_array`, covering every supported input container (numpy, list, tuple), shuffling, multi-output targets, text arrays, and the static types its overloads advertise.

river/stream/iter_array.py

@@ -2,20 +2,109 @@
 
 import itertools
 import random
+import typing
 
 import numpy as np
 
 from river import base
 
+if typing.TYPE_CHECKING:
+    from collections.abc import Callable, Iterator, Sequence
 
+Target = typing.TypeVar("Target", bound=base.typing.Target)
+"""The type of a single target value, i.e. of a row of a 1D array of targets."""
+
+Array: typing.TypeAlias = "np.ndarray | Sequence[typing.Any]"
+"""A numpy array or a plain Python sequence."""
+
+Features: typing.TypeAlias = "dict[base.typing.FeatureName, typing.Any]"
+"""One row of features, labeled, which is what an estimator takes as its `x`."""
+
+
+def _passthrough(row: str) -> str:
+    """Leaves a row as it is, which is how an array of texts is yielded."""
+    return row
+
+
+def _labeler(
+    names: Sequence[base.typing.FeatureName], first_row: typing.Any
+) -> Callable[[typing.Any], Features]:
+    """Picks the labeler for every row from the first one, all rows being of the same type."""
+    if isinstance(first_row, np.ndarray):
+
+        def label_numpy(row: np.ndarray) -> Features:
+            return dict(zip(names, row.tolist()))
+
+        return label_numpy
+
+    def label(row: typing.Any) -> Features:
+        return dict(zip(names, row))
+
+    return label
+
+
+def _take(values: Array, order: Sequence[int]) -> Array:
+    """Reorders rows, with fancy indexing for numpy and row lookups for plain sequences."""
+    return values[order] if isinstance(values, np.ndarray) else [values[i] for i in order]
+
+
+# NOTE: A row of X is labeled to make a dictionary of features, unless X is an array of texts,
+# which are yielded as they are. A row of y is a target, unless y is multi-output, in which case
+# it is a dictionary too. A plain Python sequence says which of these it is, and the overloads
+# below follow it; a numpy array hands out `Any`, which is where the trail stops. The first
+# overload overlaps the wider ones by design, which is what the pyright suppression is for.
+@typing.overload
+def iter_array(  # pyright: ignore[reportOverlappingOverload]
+    X: Sequence[str],
+    y: Sequence[Target],
+    feature_names: list[base.typing.FeatureName] | None = None,
+    target_names: list[base.typing.FeatureName] | None = None,
+    shuffle: bool = False,
+    seed: int | None = None,
+) -> Iterator[tuple[str, Target]]: ...
+
+
+@typing.overload
 def iter_array(
-    X: np.ndarray,
-    y: np.ndarray | None = None,
+    X: Sequence[str],
+    y: Array | None = None,
     feature_names: list[base.typing.FeatureName] | None = None,
     target_names: list[base.typing.FeatureName] | None = None,
     shuffle: bool = False,
     seed: int | None = None,
-) -> base.typing.Stream:
+) -> Iterator[tuple[str, typing.Any]]: ...
+
+
+@typing.overload
+def iter_array(
+    X: Array,
+    y: Sequence[Target],
+    feature_names: list[base.typing.FeatureName] | None = None,
+    target_names: list[base.typing.FeatureName] | None = None,
+    shuffle: bool = False,
+    seed: int | None = None,
+) -> Iterator[tuple[Features, Target]]: ...
+
+
+@typing.overload
+def iter_array(
+    X: Array,
+    y: Array | None = None,
+    feature_names: list[base.typing.FeatureName] | None = None,
+    target_names: list[base.typing.FeatureName] | None = None,
+    shuffle: bool = False,
+    seed: int | None = None,
+) -> base.typing.Stream: ...
+
+
+def iter_array(
+    X: Array,
+    y: Array | None = None,
+    feature_names: list[base.typing.FeatureName] | None = None,
+    target_names: list[base.typing.FeatureName] | None = None,
+    shuffle: bool = False,
+    seed: int | None = None,
+) -> Iterator[tuple[Features | str, typing.Any]]:
     """Iterates over the rows from an array of features and an array of targets.
 
     This method is intended to work with `numpy` arrays, but should also work with Python lists.
@@ -69,36 +158,31 @@ def iter_array(
     bar False
 
     """
+    if (n_rows := len(X)) == 0:
+        return
+
+    if shuffle:
+        order = random.Random(seed).sample(range(n_rows), k=n_rows)
+        X = _take(X, order)
+        y = y if y is None else _take(y, order)
 
+    handle_features: Callable[[typing.Any], Features | str]
     # If the first row of X is actually a string, then we assume all the rows are strings and will
-    # pass them through
+    # pass them through. If not we assume each row is a set of features, and will label them.
     if isinstance(X[0], str):
-
-        def handle_features(x):
-            return x.tolist() if isinstance(x, np.ndarray) else x
-
-    # If not we assume each row if a set of features, and will convert them to a dictionary
+        handle_features = _passthrough
     else:
-        feature_names = list(range(len(X[0]))) if feature_names is None else feature_names
-
-        def handle_features(x):
-            return dict(zip(feature_names, xi.tolist() if isinstance(xi, np.ndarray) else xi))
-
-    multioutput = y is not None and not np.isscalar(y[0])
-    if multioutput and target_names is None:
-        target_names = list(range(len(y[0])))  # type: ignore
-
-    # Shuffle the data
-    rng = random.Random(seed)
-    if shuffle:
-        order = rng.sample(range(len(X)), k=len(X))
-        X = X[order]
-        y = y if y is None else y[order]
+        handle_features = _labeler(
+            range(len(X[0])) if feature_names is None else feature_names, X[0]
+        )
 
-    if multioutput:
-        for xi, yi in itertools.zip_longest(X, y if hasattr(y, "__iter__") else []):  # type: ignore
-            yield handle_features(xi), dict(zip(target_names, yi.tolist()))  # type: ignore
+    # zip_longest pads a target array shorter than X with Nones, instead of stopping short of it.
+    rows = itertools.zip_longest(X, () if y is None else y)
 
+    if y is not None and not np.isscalar(y[0]):
+        handle_target = _labeler(range(len(y[0])) if target_names is None else target_names, y[0])
+        for xi, yi in rows:
+            yield handle_features(xi), handle_target(yi)
     else:
-        for xi, yi in itertools.zip_longest(X, y if hasattr(y, "__iter__") else []):  # type: ignore
+        for xi, yi in rows:
             yield handle_features(xi), yi.item() if isinstance(yi, np.generic) else yi

tests/stream/test_iter_array.py

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+import operator
+import typing
+
+import numpy as np
+import pytest
+
+from river import stream
+
+# Imported at runtime, rather than under TYPE_CHECKING, as `assert_type` evaluates its arguments.
+from river.base.typing import FeatureName
+
+if typing.TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+    Array = np.ndarray | Sequence[typing.Any]
+    Rows = list[tuple[typing.Any, typing.Any]]
+
+# The same data reaches `iter_array` as numpy or as plain Python containers, and both are
+# documented as supported. Every case below is expected to yield the very same stream.
+ARRAY_BACKENDS: dict[str, Callable[[list[typing.Any]], Array]] = {
+    "numpy": np.array,
+    "list": list,
+    "tuple": tuple,
+}
+
+FEATURES = [[1, 2, 3], [11, 12, 13]]
+LABELED = [{0: 1, 1: 2, 2: 3}, {0: 11, 1: 12, 2: 13}]
+TARGET = [True, False]
+MULTI_TARGET = [[1, 2], [11, 12]]
+TEXTS = ["foo", "bar"]
+
+# Enough rows for a shuffle to be all but certain to reorder them, each labeled with its own
+# first feature so that a row and its target can be checked to travel together.
+ROWS = [[i, i * 10] for i in range(10)]
+LABELS = list(range(10))
+MULTI_LABELS = [[i, -i] for i in range(10)]
+LONG_TEXTS = [f"text {i}" for i in range(10)]
+SHUFFLE: dict[str, typing.Any] = {"shuffle": True, "seed": 42}
+
+
+@pytest.fixture(params=list(ARRAY_BACKENDS))
+def array(request: pytest.FixtureRequest) -> Callable[[list[typing.Any]], Array]:
+    """Yields one array-like constructor per supported input container."""
+    return ARRAY_BACKENDS[request.param]
+
+
+class Case(typing.NamedTuple):
+    """One call to `iter_array`, with its arrays expressed as plain Python containers."""
+
+    X: list[typing.Any]
+    y: list[typing.Any] | None = None
+    kwargs: dict[str, typing.Any] = {}
+
+
+def run(case: Case, backend: Callable[[list[typing.Any]], Array]) -> Rows:
+    return list(
+        stream.iter_array(
+            backend(case.X), None if case.y is None else backend(case.y), **case.kwargs
+        )
+    )
+
+
+STREAMS: dict[str, tuple[Case, Rows]] = {
+    "features-only": (Case(FEATURES), [(LABELED[0], None), (LABELED[1], None)]),
+    "with-target": (Case(FEATURES, TARGET), [(LABELED[0], True), (LABELED[1], False)]),
+    "named-features": (
+        Case(FEATURES, TARGET, {"feature_names": ["x1", "x2", "x3"]}),
+        [({"x1": 1, "x2": 2, "x3": 3}, True), ({"x1": 11, "x2": 12, "x3": 13}, False)],
+    ),
+    "fewer-names-than-features": (
+        Case(FEATURES, kwargs={"feature_names": ["x1"]}),
+        [({"x1": 1}, None), ({"x1": 11}, None)],
+    ),
+    "shorter-target-is-padded": (
+        Case(FEATURES, [True]),
+        [(LABELED[0], True), (LABELED[1], None)],
+    ),
+    "multioutput": (
+        Case(FEATURES, MULTI_TARGET),
+        [(LABELED[0], {0: 1, 1: 2}), (LABELED[1], {0: 11, 1: 12})],
+    ),
+    "named-outputs": (
+        Case(FEATURES, MULTI_TARGET, {"target_names": ["y1", "y2"]}),
+        [(LABELED[0], {"y1": 1, "y2": 2}), (LABELED[1], {"y1": 11, "y2": 12})],
+    ),
+    "target-names-ignored-for-a-single-output": (
+        Case(FEATURES, TARGET, {"target_names": ["y1"]}),
+        [(LABELED[0], True), (LABELED[1], False)],
+    ),
+    "text-passes-through": (Case(TEXTS, TARGET), [("foo", True), ("bar", False)]),
+    "empty": (Case([], []), []),
+    "empty-without-target": (Case([]), []),
+}
+
+
+@pytest.mark.parametrize(("case", "expected"), STREAMS.values(), ids=STREAMS)
+def test_expected_stream(
+    case: Case, expected: Rows, array: Callable[[list[typing.Any]], Array]
+) -> None:
+    """Each input shape yields its expected rows, whichever container the arrays come in.
+
+    Features are labeled with their position when no names are given, and features without a
+    name are dropped. When `y` is omitted, or shorter than `X`, the target is padded with `None`.
+    A 2D target yields one dict per row, and a 1D array of texts is yielded as-is.
+    """
+    assert run(case, array) == expected
+
+
+SHUFFLED_STREAMS: dict[str, Case] = {
+    "with-target": Case(ROWS, LABELS, SHUFFLE),
+    "multioutput": Case(ROWS, MULTI_LABELS, SHUFFLE),
+    "without-target": Case(ROWS, None, SHUFFLE),
+    "text": Case(LONG_TEXTS, LABELS, SHUFFLE),
+}
+
+
+@pytest.mark.parametrize("case", SHUFFLED_STREAMS.values(), ids=SHUFFLED_STREAMS)
+def test_backends_agree_on_shuffling(case: Case) -> None:
+    """numpy, lists and tuples are all reordered the same way for a given seed."""
+    streams = [run(case, backend) for backend in ARRAY_BACKENDS.values()]
+    assert all(rows == streams[0] for rows in streams)
+
+
+def test_shuffle_reorders_and_preserves_rows(array: Callable[[list[typing.Any]], Array]) -> None:
+    """Shuffling only reorders the rows, each one keeping its own target."""
+    X, y = array(ROWS), array(LABELS)
+    plain = list(stream.iter_array(X, y))
+    shuffled = list(stream.iter_array(X, y, shuffle=True, seed=42))
+
+    assert shuffled != plain
+    assert all(xi[0] == yi for xi, yi in shuffled)
+    assert sorted(shuffled, key=operator.itemgetter(1)) == plain
+
+
+def test_shuffle_is_seeded(array: Callable[[list[typing.Any]], Array]) -> None:
+    """The same seed always gives the same order, and different seeds give different ones."""
+    X, y = array(ROWS), array(LABELS)
+
+    assert list(stream.iter_array(X, y, **SHUFFLE)) == list(stream.iter_array(X, y, **SHUFFLE))
+    assert list(stream.iter_array(X, y, shuffle=True, seed=0)) != list(
+        stream.iter_array(X, y, shuffle=True, seed=1)
+    )
+
+
+def test_native_python_scalars(array: Callable[[list[typing.Any]], Array]) -> None:
+    """Cells and targets are plain Python values, not numpy scalars."""
+    xi, yi = next(iter(stream.iter_array(array(FEATURES), array(TARGET))))
+    _, multi = next(iter(stream.iter_array(array(FEATURES), array(MULTI_TARGET))))
+    text, _ = next(iter(stream.iter_array(array(TEXTS), array(TARGET))))
+
+    assert [type(value) for value in xi.values()] == [int, int, int]
+    assert type(yi) is bool
+    assert [type(value) for value in multi.values()] == [int, int]
+    assert isinstance(text, str)
+
+
+def test_static_types_follow_the_input_shapes() -> None:
+    """What a checker makes of each input shape, which `assert_type` fails the mypy run over."""
+    # A sequence of texts yields texts, and a sequence of targets yields the target's own type.
+    typing.assert_type(next(iter(stream.iter_array(["a", "b"], [1, 2]))), tuple[str, int])
+    typing.assert_type(
+        next(iter(stream.iter_array([[1, 2]], [True]))),
+        tuple[dict[FeatureName, typing.Any], bool],
+    )
+    # numpy hands out `Any`, and so does a target a checker cannot tell apart from a 2D one.
+    typing.assert_type(
+        next(iter(stream.iter_array(np.array([[1, 2]]), np.array([1])))),
+        tuple[dict[FeatureName, typing.Any], typing.Any],
+    )
+    typing.assert_type(
+        next(iter(stream.iter_array([[1, 2]], [[1, 2]]))),
+        tuple[dict[FeatureName, typing.Any], typing.Any],
+    )
+
+
+def test_iteration_is_lazy() -> None:
+    """Nothing is read from the arrays until the stream is iterated over."""
+    dataset = stream.iter_array(np.array([1, 2, 3]))
+    with pytest.raises(TypeError):
+        _ = next(iter(dataset))