diff --git a/.gitignore b/.gitignore index 16ac5e5..de0c386 100644 --- a/.gitignore +++ b/.gitignore @@ -13,4 +13,5 @@ coverage.info .coverage htmlcov/ -.claude \ No newline at end of file +.claude +.agents \ No newline at end of file diff --git a/README.md b/README.md index 119e37e..9b918ca 100644 --- a/README.md +++ b/README.md @@ -93,12 +93,12 @@ Note: Most `Itr` methods are **lazy transformations**, meaning they return a new `Itr` instance without immediately processing any data. This allows for arbitrary chaining and efficient memory usage, as items are only processed as they are requested. In most cases, `Itr` simply acts as a convenient wrapper around `itertools`, enabling this left-to-right chaining syntax. - **Combining and splitting:** `partition`, `copy`, `batched`, `pairwise`, `rolling`, `chain`, `cycle`, `repeat`, `product`, `inspect`, `intersperse`, `interleave`, `chunk_by`, `zip_longest` -- **Transformation and filtering:** `accumulate`, `filter`, `map`, `starmap`, `map_while`, `flatten`, `flat_map`, `skip_while`, `take_while`, `dedup` +- **Transformation and filtering:** `accumulate`, `filter`, `map`, `starmap`, `map_while`, `flatten`, `flat_map`, `skip_while`, `take_while`, `dedup`, `dedup_with_count`, `scan` However, some methods are **eager consumers**. These methods iterate over and consume the underlying data, returning concrete values, collections, or aggregates. Examples include: * **Collection methods:** `collect`, `last`, `next`, `next_chunk`, `next_if`, `nth`, `position` -* **Aggregation methods:** `count`, `reduce`, `max`, `min`, `sum`, `prod`, `all`, `any`, `consume`, `find`, `fold` +* **Aggregation methods:** `count`, `reduce`, `max`, `min`, `sum`, `prod`, `all`, `any`, `consume`, `find`, `fold`, `eq`, `is_sorted` * **Sorting/grouping:** `sorted_by` and `groupby` sort the entire input up front, and `value_counts` counts it (most common first, like pandas), so all three consume the whole iterator immediately and must not be used on infinite sources. Use the lazy `chunk_by` to group consecutive runs without sorting. ### Important Considerations diff --git a/doc/apidoc.md b/doc/apidoc.md index 85b9a6b..4cd2aa9 100644 --- a/doc/apidoc.md +++ b/doc/apidoc.md @@ -182,6 +182,23 @@ Example: (1, 2, 3, 1) +### `dedup_with_count` + +Lazily collapse each *consecutive* run of equal items into a (item, count) pair (run-length encoding). + +The lazy, positional counterpart to `value_counts`: this counts adjacent runs and preserves order (so the +same item may appear more than once), where `value_counts` counts occurrences over the whole iterator and is +eager. Items are compared by equality and do not need to be hashable. Works on infinite iterators, provided +no individual run is infinite. + +Returns: + Itr[tuple[T, int]]: An iterator of (item, run length) pairs. + +Example: + >>> Itr([4, 4, 2, 3, 3, 1]).dedup_with_count().collect() + ((4, 2), (2, 1), (3, 2), (1, 1)) + + ### `enumerate` Yield pairs of (index, item) for each item in the iterator, where index starts at 0 or the value provided @@ -191,6 +208,28 @@ Returns: +### `eq` + +Compare the remaining items with another iterable, element by element (like Rust's `Iterator::eq`). + +Returns True only if both yield equal items in the same order and have the same length. Comparison +short-circuits at the first difference, so unlike `tuple(a) == tuple(b)` neither side is fully materialised. +NB This consumes as much of the iterator as it needs to. Note that `Itr` does not define `__eq__`, so `==` +compares identity, not contents. + +Args: + other (Iterable[Any]): The iterable to compare against. + +Returns: + bool: True if the sequences are element-wise equal. + +Example: + >>> Itr([1, 2, 3]).eq([1, 2, 3]) + True + >>> Itr([1, 2, 3]).eq([1, 2]) + False + + ### `filter` Yield only items that satisfy the predicate. @@ -330,6 +369,28 @@ Returns: +### `is_sorted` + +Check whether the remaining items are in sorted order (like Rust's `Iterator::is_sorted`). + +Order is non-strict, so runs of equal items are sorted. An empty or single-item iterator is sorted. The +check short-circuits at the first item out of order, but NB it consumes the iterator either way. + +Args: + key (Callable[[T], Any] | None): Applied to each item before comparison, as in `sorted_by` (covering + Rust's `is_sorted_by_key`). Defaults to comparing the items themselves. + reverse (bool): If True, check for descending rather than ascending order. + +Returns: + bool: True if the items are in the expected order. + +Example: + >>> Itr([1, 2, 2, 3]).is_sorted() + True + >>> Itr(["ccc", "bb", "a"]).is_sorted(len, reverse=True) + True + + ### `last` Return the last item from the iterator. Do not use on an open-ended Iterable @@ -609,6 +670,29 @@ Rolling window (generalisation of pairwise) Rather than copying the iterator multiple times, collect n, yield the sequence and incrementally drop/add +### `scan` + +Lazily map items through a running state, optionally stopping early (like Rust's `Iterator::scan`). + +`func` receives the current state and the next item, and returns either a `(new_state, output)` pair or +None to stop iterating. This generalises `accumulate`: the state need not be the same type as the items, +and iteration can terminate on a condition. Yielding None as an *output* is unambiguous, since the halt +signal is the entire return value rather than the output value. + +Args: + init (S): The initial state. + func (Callable[[S, T], tuple[S, U] | None]): Maps (state, item) to (new state, output), or None to stop. + +Returns: + Itr[U]: An iterator over the outputs. + +Example: + >>> Itr([1, 2, 3, 4]).scan(0, lambda total, x: (total + x, total + x)).collect() + (1, 3, 6, 10) + >>> Itr([1, 2, 3, 4]).scan(0, lambda total, x: None if total + x > 5 else (total + x, total + x)).collect() + (1, 3) + + ### `skip` Skip the next n items in the iterator. diff --git a/relnotes.md b/relnotes.md index 3496152..f05463b 100644 --- a/relnotes.md +++ b/relnotes.md @@ -3,6 +3,10 @@ ### New features - Installable **agent skill**: the package now bundles a `SKILL.md` reference for AI coding agents, plus an `itrx-skill` console script to symlink it into a project (`itrx-skill --install [PATH]` / `--remove [PATH]`, default `PATH=.agents`, creating `PATH/skills/itrx`). The symlink points at the skill inside the installed `itrx`, so it always matches the version in use. See the "Agent skill" section of the README. +- `dedup_with_count()`: the lazy, positional counterpart to `value_counts` — collapses each *consecutive* run of equal items into an `(item, count)` pair (run-length encoding), preserving order and working on infinite iterators. Note the `(item, count)` ordering matches `value_counts` and is the reverse of Rust's `dedup_with_count`. +- `scan(init, func)`: lazily map items through a running state, optionally stopping early (Rust's `Iterator::scan`). Generalises `accumulate`: the state need not share the items' type, and returning `None` halts iteration. Returning `(new_state, None)` still yields `None` as an output, so the halt signal is never ambiguous. +- `is_sorted(key=None, *, reverse=False)`: check whether the remaining items are in order, short-circuiting at the first inversion. Non-strict, so equal runs count as sorted. Covers Rust's `is_sorted` and `is_sorted_by_key`, and adds `reverse` for symmetry with `sorted_by`. +- `eq(other)`: element-wise comparison against another iterable, short-circuiting at the first difference instead of materialising both sides (Rust's `Iterator::eq`). Note `Itr` defines no `__eq__`, so `==` remains an identity check. ## 0.4.0 diff --git a/src/itrx/itr.py b/src/itrx/itr.py index b56f243..596abb5 100644 --- a/src/itrx/itr.py +++ b/src/itrx/itr.py @@ -186,6 +186,24 @@ def dedup(self) -> "Itr[T]": """ return Itr(k for k, _ in itertools.groupby(self._it)) + def dedup_with_count(self) -> "Itr[tuple[T, int]]": + """Lazily collapse each *consecutive* run of equal items into a (item, count) pair (run-length encoding). + + The lazy, positional counterpart to `value_counts`: this counts adjacent runs and preserves order (so the + same item may appear more than once), where `value_counts` counts occurrences over the whole iterator and is + eager. Items are compared by equality and do not need to be hashable. Works on infinite iterators, provided + no individual run is infinite. + + Returns: + Itr[tuple[T, int]]: An iterator of (item, run length) pairs. + + Example: + >>> Itr([4, 4, 2, 3, 3, 1]).dedup_with_count().collect() + ((4, 2), (2, 1), (3, 2), (1, 1)) + """ + # note the (item, count) ordering matches value_counts, and is the reverse of Rust's dedup_with_count + return cast("Itr[tuple[T, int]]", Itr((k, sum(1 for _ in g)) for k, g in itertools.groupby(self._it))) + def enumerate(self, *, start: int = 0) -> "Itr[tuple[int, T]]": """Yield pairs of (index, item) for each item in the iterator, where index starts at 0 or the value provided @@ -195,6 +213,29 @@ def enumerate(self, *, start: int = 0) -> "Itr[tuple[int, T]]": """ return cast("Itr[tuple[int, T]]", Itr(enumerate(self._it, start))) + def eq(self, other: Iterable[Any]) -> bool: + """Compare the remaining items with another iterable, element by element (like Rust's `Iterator::eq`). + + Returns True only if both yield equal items in the same order and have the same length. Comparison + short-circuits at the first difference, so unlike `tuple(a) == tuple(b)` neither side is fully materialised. + NB This consumes as much of the iterator as it needs to. Note that `Itr` does not define `__eq__`, so `==` + compares identity, not contents. + + Args: + other (Iterable[Any]): The iterable to compare against. + + Returns: + bool: True if the sequences are element-wise equal. + + Example: + >>> Itr([1, 2, 3]).eq([1, 2, 3]) + True + >>> Itr([1, 2, 3]).eq([1, 2]) + False + """ + unequal = object() + return all(a == b for a, b in itertools.zip_longest(self._it, other, fillvalue=unequal)) + def filter(self, predicate: Predicate[T]) -> "Itr[T]": """Yield only items that satisfy the predicate. @@ -390,6 +431,31 @@ def interleaver() -> Generator[T | U, None, None]: return cast("Itr[T | U]", Itr(interleaver())) + def is_sorted(self, key: Callable[[T], Any] | None = None, *, reverse: bool = False) -> bool: + """Check whether the remaining items are in sorted order (like Rust's `Iterator::is_sorted`). + + Order is non-strict, so runs of equal items are sorted. An empty or single-item iterator is sorted. The + check short-circuits at the first item out of order, but NB it consumes the iterator either way. + + Args: + key (Callable[[T], Any] | None): Applied to each item before comparison, as in `sorted_by` (covering + Rust's `is_sorted_by_key`). Defaults to comparing the items themselves. + reverse (bool): If True, check for descending rather than ascending order. + + Returns: + bool: True if the items are in the expected order. + + Example: + >>> Itr([1, 2, 2, 3]).is_sorted() + True + >>> Itr(["ccc", "bb", "a"]).is_sorted(len, reverse=True) + True + """ + # T is unbounded so is not known to be orderable, as in sorted_by/groupby + keyed = cast("Iterable[Any]", self._it if key is None else (key(item) for item in self._it)) + pairs = itertools.pairwise(keyed) + return all(b <= a for a, b in pairs) if reverse else all(a <= b for a, b in pairs) + def last(self) -> T: """Return the last item from the iterator. Do not use on an open-ended Iterable @@ -686,6 +752,39 @@ def rolling(self, n: int) -> "Itr[tuple[T, ...]]": shifted_iterators = (itertools.islice(it, i, None) for i, it in enumerate(iterators)) return cast("Itr[tuple[T, ...]]", Itr(zip(*shifted_iterators, strict=False))) + def scan[S, U](self, init: S, func: Callable[[S, T], tuple[S, U] | None]) -> "Itr[U]": + """Lazily map items through a running state, optionally stopping early (like Rust's `Iterator::scan`). + + `func` receives the current state and the next item, and returns either a `(new_state, output)` pair or + None to stop iterating. This generalises `accumulate`: the state need not be the same type as the items, + and iteration can terminate on a condition. Yielding None as an *output* is unambiguous, since the halt + signal is the entire return value rather than the output value. + + Args: + init (S): The initial state. + func (Callable[[S, T], tuple[S, U] | None]): Maps (state, item) to (new state, output), or None to stop. + + Returns: + Itr[U]: An iterator over the outputs. + + Example: + >>> Itr([1, 2, 3, 4]).scan(0, lambda total, x: (total + x, total + x)).collect() + (1, 3, 6, 10) + >>> Itr([1, 2, 3, 4]).scan(0, lambda total, x: None if total + x > 5 else (total + x, total + x)).collect() + (1, 3) + """ + + def gen() -> Generator[U, None, None]: + state = init + for item in self._it: + result = func(state, item) + if result is None: + return + state, output = result + yield output + + return Itr(gen()) + def skip(self, n: int) -> "Itr[T]": """Skip the next n items in the iterator. diff --git a/src/itrx/skill/SKILL.md b/src/itrx/skill/SKILL.md index d80de39..ba638cc 100644 --- a/src/itrx/skill/SKILL.md +++ b/src/itrx/skill/SKILL.md @@ -48,8 +48,8 @@ equivalent `itertools` code, because it *is* that code underneath. Reach for it new `Itr` and pulls items only on demand, so an infinite source stays workable right up to the terminal call. - **The operation exists in Rust's `Iterator` but not as a Python builtin** — `fold`, `inspect`, - `partition`, `position`, `intersperse`, `interleave`, `dedup`, `chunk_by`, `unzip`, `map_while`, - `next_chunk`, `step_by`, `rolling`. + `partition`, `position`, `intersperse`, `interleave`, `dedup`, `dedup_with_count`, `chunk_by`, + `unzip`, `map_while`, `scan`, `is_sorted`, `eq`, `next_chunk`, `step_by`, `rolling`. Conversely, it is **not** worth it for a single `map`/`filter` (a comprehension is clearer), for code already dominated by numpy/pandas vectorised calls, or where the data is a materialised @@ -67,22 +67,25 @@ wrapping a **generator or iterator** hands over ownership: consuming the `Itr` c sources — with the exception of `product`, which materialises `other` up front like `itertools.product` does): -`accumulate`, `batched`, `chain`, `chunk_by`, `copy`, `cycle`, `dedup`, `enumerate`, `filter`, +`accumulate`, `batched`, `chain`, `chunk_by`, `copy`, `cycle`, `dedup`, `dedup_with_count`, +`enumerate`, `filter`, `flat_map`, `flatten`, `inspect`, `interleave`, `intersperse`, `map`, `map_dict`, `map_while`, -`pairwise`, `partition`, `product`, `repeat`, `rolling`, `skip`, `skip_while`, `step_by`, `take`, +`pairwise`, `partition`, `product`, `repeat`, `rolling`, `scan`, `skip`, `skip_while`, +`step_by`, `take`, `take_while`, `tee`, `unzip`, `zip`, `zip_longest` **Eager** (consume the iterator, return a concrete value; **never** on an infinite source): - Collection: `collect`, `last`, `next`, `next_chunk`, `next_if`, `nth`, `position`, `peek` -- Aggregation: `all`, `any`, `consume`, `count`, `find`, `fold`, `for_each`, `max`, `min`, `prod`, - `reduce`, `sum` +- Aggregation: `all`, `any`, `consume`, `count`, `eq`, `find`, `fold`, `for_each`, `is_sorted`, + `max`, `min`, `prod`, `reduce`, `sum` - Whole-input reordering: `groupby`, `sorted_by`, `value_counts`, `rev` Note that some of these consume only as far as they need to: `next`, `next_chunk`, `nth`, `next_if`, `peek`, `find`, `position`, `any` and `all` short-circuit, so they *are* safe on an infinite source. `collect`, `count`, `last`, `consume`, `fold`, `reduce`, `sum`, `prod`, `max`, -`min`, `for_each`, `rev`, `groupby`, `sorted_by` and `value_counts` are not. +`min`, `for_each`, `rev`, `groupby`, `sorted_by` and `value_counts` are not. `eq` and `is_sorted` +short-circuit on the first difference or inversion, so they too are safe on an infinite source. ## Outputs @@ -126,6 +129,14 @@ infinite source. `collect`, `count`, `last`, `consume`, `fold`, `reduce`, `sum`, - **`dedup()` removes only *adjacent* duplicates**, keeping the first of each run. It compares by equality (items need not be hashable) and stays lazy — it is not "unique". For global uniqueness use `collect(set)`, accepting the loss of order. +- **`dedup_with_count()` is run-length encoding** — the same adjacent-run logic as `dedup`, but + yielding `(item, count)` pairs. It is the lazy, positional counterpart to `value_counts`: same + output shape, but counting adjacent runs in source order rather than occurrences overall. On + `[4, 4, 2, 3, 3, 1]` it gives `((4, 2), (2, 1), (3, 2), (1, 1))` where `value_counts()` gives + `((4, 2), (3, 2), (2, 1), (1, 1))`. Prefer it to + `chunk_by(f).map(lambda kv: (kv[0], len(kv[1])))`, which materialises every run to measure it. + Note the `(item, count)` order is the reverse of Rust's `dedup_with_count`. It stays lazy on an + infinite source, but an infinite individual run (e.g. `itertools.repeat(1)`) will hang. - **`rev()` materialises the entire remaining sequence** into memory before yielding — unavoidable, but never call it on an unbounded source. - **`repeat(n)` tees the iterator `n` times**, so it buffers the whole sequence for large `n`; it @@ -149,6 +160,15 @@ infinite source. `collect`, `count`, `last`, `consume`, `fold`, `reduce`, `sum`, must be finite even though the chain stays lazy in `self`. - **`inspect(func)` is the lazy debugging hook** — it calls `func` on each item and passes it through unchanged, so you can drop it mid-chain without altering results. +- **`scan(init, func)` is `accumulate` with a separate state type and an early exit.** `func(state, + item)` returns `(new_state, output)`, or `None` to stop. `None` as the *whole* return value halts; + `(new_state, None)` yields `None` as an output, so the two are never ambiguous. Reach for + `accumulate` when the state is just the running value, and `scan` otherwise. +- **`is_sorted(key=None, *, reverse=False)` is non-strict** — runs of equal items count as sorted, + and empty/single-item iterators are sorted. The `key` argument covers Rust's `is_sorted_by_key`. +- **`eq(other)` compares contents; `==` does not.** `Itr` defines no `__eq__`, so `itr_a == itr_b` + is an identity check. Use `.eq(other)` for an element-wise comparison, which also short-circuits + rather than materialising both sides. ## Typing diff --git a/src/test/test_aggregation.py b/src/test/test_aggregation.py index 5848fe1..7610463 100644 --- a/src/test/test_aggregation.py +++ b/src/test/test_aggregation.py @@ -1,3 +1,4 @@ +import itertools from operator import add, mul, sub, truediv import pytest @@ -214,3 +215,58 @@ def test_prod_empty() -> None: def test_prod_with_zero() -> None: assert Itr([1, 0, 5]).prod() == 0 + + +def test_eq() -> None: + assert Itr([1, 2, 3]).eq([1, 2, 3]) + assert not Itr([1, 2, 3]).eq([1, 2, 4]) + + +def test_eq_empty() -> None: + assert Itr[int]([]).eq([]) + assert not Itr[int]([]).eq([1]) + + +def test_eq_differing_lengths() -> None: + assert not Itr([1, 2, 3]).eq([1, 2]) + assert not Itr([1, 2]).eq([1, 2, 3]) + + +def test_eq_accepts_any_iterable() -> None: + assert Itr("abc").eq(iter("abc")) + assert Itr([1, 2, 3]).eq(Itr(range(1, 4))) + + +def test_eq_short_circuits_on_infinite() -> None: + # differs at the third item, so neither side is exhausted + assert not Itr(itertools.count()).eq(itertools.chain([0, 1], [99])) + + +def test_is_sorted() -> None: + assert Itr([1, 2, 3]).is_sorted() + assert not Itr([1, 3, 2]).is_sorted() + + +def test_is_sorted_allows_equal_items() -> None: + assert Itr([1, 2, 2, 3]).is_sorted() + + +def test_is_sorted_empty_and_single() -> None: + assert Itr[int]([]).is_sorted() + assert Itr([5]).is_sorted() + + +def test_is_sorted_with_key() -> None: + assert Itr(["a", "bb", "ccc"]).is_sorted(len) + assert not Itr(["ccc", "a"]).is_sorted(len) + + +def test_is_sorted_reverse() -> None: + assert Itr([3, 2, 1]).is_sorted(reverse=True) + assert not Itr([1, 2, 3]).is_sorted(reverse=True) + assert Itr(["ccc", "bb", "a"]).is_sorted(len, reverse=True) + + +def test_is_sorted_agrees_with_sorted_by() -> None: + data = ["ccc", "a", "bb"] + assert Itr(data).sorted_by(len).is_sorted(len) diff --git a/src/test/test_transform_filter.py b/src/test/test_transform_filter.py index 61bd33a..dca3004 100644 --- a/src/test/test_transform_filter.py +++ b/src/test/test_transform_filter.py @@ -229,6 +229,44 @@ def test_dedup_lazy_on_infinite() -> None: assert it.take(3).collect() == (0, 1, 2) +def test_dedup_with_count() -> None: + assert Itr([4, 4, 2, 3, 3, 1]).dedup_with_count().collect() == ((4, 2), (2, 1), (3, 2), (1, 1)) + + +def test_dedup_with_count_no_duplicates() -> None: + assert Itr([1, 2, 3]).dedup_with_count().collect() == ((1, 1), (2, 1), (3, 1)) + + +def test_dedup_with_count_empty() -> None: + assert Itr[int]([]).dedup_with_count().collect() == () + + +def test_dedup_with_count_single_run() -> None: + assert Itr([7, 7, 7, 7]).dedup_with_count().collect() == ((7, 4),) + + +def test_dedup_with_count_unhashable_items() -> None: + assert Itr([[1], [1], [2]]).dedup_with_count().collect() == (([1], 2), ([2], 1)) + + +def test_dedup_with_count_lazy_on_infinite() -> None: + it = Itr(itertools.count()).flat_map(lambda x: (x, x)).dedup_with_count() + assert it.take(3).collect() == ((0, 2), (1, 2), (2, 2)) + + +def test_dedup_with_count_keys_match_dedup() -> None: + data = [1, 1, 2, 2, 2, 3, 1] + counted = Itr(data).dedup_with_count().starmap(lambda item, _n: item).collect() + assert counted == Itr(data).dedup().collect() + + +def test_dedup_with_count_differs_from_value_counts() -> None: + # dedup_with_count counts adjacent runs positionally, value_counts counts occurrences overall + data = [4, 4, 2, 3, 3, 1] + assert Itr(data).dedup_with_count().collect() == ((4, 2), (2, 1), (3, 2), (1, 1)) + assert Itr(data).value_counts().collect() == ((4, 2), (3, 2), (2, 1), (1, 1)) + + def test_sorted_by() -> None: assert Itr(["ccc", "a", "bb"]).sorted_by(len).collect() == ("a", "bb", "ccc") @@ -244,3 +282,32 @@ def test_sorted_by_is_stable() -> None: def test_sorted_by_empty() -> None: assert Itr[int]([]).sorted_by(lambda x: x).collect() == () + + +def test_scan() -> None: + assert Itr([1, 2, 3, 4]).scan(0, lambda total, x: (total + x, total + x)).collect() == (1, 3, 6, 10) + + +def test_scan_stops_on_none() -> None: + result = Itr([1, 2, 3, 4]).scan(0, lambda total, x: None if total + x > 5 else (total + x, total + x)) + assert result.collect() == (1, 3) + + +def test_scan_empty() -> None: + assert Itr[int]([]).scan(0, lambda total, x: (total + x, total + x)).collect() == () + + +def test_scan_state_type_differs_from_output() -> None: + # state is a running count, output is the item tagged with its 1-based position + result = Itr("abc").scan(0, lambda n, c: (n + 1, f"{n + 1}{c}")) + assert result.collect() == ("1a", "2b", "3c") + + +def test_scan_can_yield_none_outputs() -> None: + # None is only a halt signal as the whole return value, not as an output + assert Itr([1, 2]).scan(0, lambda total, x: (total + x, None)).collect() == (None, None) + + +def test_scan_lazy_on_infinite() -> None: + running = Itr(itertools.count(1)).scan(0, lambda total, x: (total + x, total + x)) + assert running.take(4).collect() == (1, 3, 6, 10)