Cleanup, by applying more in-built python functionality for iterators

Author: simw

src/pipedata/core/chain.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from itertools import islice
 from typing import (
     Any,
     Callable,
@@ -16,16 +16,13 @@
     overload,
 )
 
-from pipedata.core.itertools import take_next, take_up_to_n
-
 TStart = TypeVar("TStart")
 TEnd = TypeVar("TEnd")
 TOther = TypeVar("TOther")
 
 
 def _identity(input_iterator: Iterator[TEnd]) -> Iterator[TEnd]:
-    while (element := take_next(input_iterator)) is not None:
-        yield element
+    yield from input_iterator
 
 
 class CountingIterator(Iterator[TStart]):
@@ -48,38 +45,31 @@ def get_count(self) -> int:
         return self._count
 
 
-class CountedFunc(Generic[TStart, TEnd]):
+class ChainLink(Generic[TStart, TEnd]):
     def __init__(
         self,
         func: Callable[[Iterator[TStart]], Iterator[TEnd]],
     ) -> None:
         self._func = func
-        self._counting_input: Optional[CountingIterator[TStart]] = None
-        self._counting_output: Optional[CountingIterator[TEnd]] = None
+        self._input: Optional[CountingIterator[TStart]] = None
+        self._output: Optional[CountingIterator[TEnd]] = None
 
     @property
     def __name__(self) -> str:  # noqa: A003
         return self._func.__name__
 
     def __call__(self, input_iterator: Iterator[TStart]) -> Iterator[TEnd]:
-        self._counting_input = CountingIterator(input_iterator)
-        self._counting_output = CountingIterator(self._func(self._counting_input))
-        return self._counting_output
+        self._input = CountingIterator(input_iterator)
+        self._output = CountingIterator(self._func(self._input))
+        return self._output
 
     def get_counts(self) -> Tuple[int, int]:
         return (
-            0 if self._counting_input is None else self._counting_input.get_count(),
-            0 if self._counting_output is None else self._counting_output.get_count(),
+            0 if self._input is None else self._input.get_count(),
+            0 if self._output is None else self._output.get_count(),
         )
 
 
-@dataclass
-class StepCount:
-    name: str
-    inputs: int
-    outputs: int
-
-
 class Chain(Generic[TStart, TEnd]):
     @overload
     def __init__(
@@ -106,11 +96,11 @@ def __init__(
         ],
     ) -> None:
         self._previous_steps = previous_steps
-        self._func = CountedFunc(func)
+        self._func = ChainLink(func)
 
     def __call__(self, input_iterator: Iterator[TStart]) -> Iterator[TEnd]:
         if self._previous_steps is None:
-            func = cast(CountedFunc[TStart, TEnd], self._func)
+            func = cast(ChainLink[TStart, TEnd], self._func)
             return func(input_iterator)
 
         return self._func(self._previous_steps(input_iterator))  # type: ignore
@@ -133,9 +123,7 @@ def filter(self, func: Callable[[TEnd], bool]) -> Chain[TStart, TEnd]:  # noqa:
         """
 
         def new_action(previous_step: Iterator[TEnd]) -> Iterator[TEnd]:
-            while (element := take_next(previous_step)) is not None:
-                if func(element) is True:
-                    yield element
+            return filter(func, previous_step)
 
         new_action.__name__ = func.__name__
         return self.flat_map(new_action)
@@ -148,8 +136,7 @@ def map(  # noqa: A003
         """
 
         def new_action(previous_step: Iterator[TEnd]) -> Iterator[TOther]:
-            while (element := take_next(previous_step)) is not None:
-                yield func(element)
+            return map(func, previous_step)
 
         new_action.__name__ = func.__name__
         return self.flat_map(new_action)
@@ -165,7 +152,7 @@ def map_tuple(
         """
 
         def new_action(previous_step: Iterator[TEnd]) -> Iterator[TOther]:
-            while elements := take_up_to_n(previous_step, n):
+            while elements := tuple(islice(previous_step, n)):
                 yield func(elements)
 
         new_action.__name__ = func.__name__

src/pipedata/core/itertools.py

@@ -1,7 +1,6 @@
 from typing import Iterator, Tuple
 
 from pipedata.core import Chain, ChainStart
-from pipedata.core.itertools import take_next
 
 
 def test_chain() -> None:
@@ -86,7 +85,7 @@ def is_even(value: int) -> bool:
 
 def test_chain_flat_map() -> None:
     def add_one(input_iterator: Iterator[int]) -> Iterator[int]:
-        while (element := take_next(input_iterator)) is not None:
+        for element in input_iterator:
             yield element + 1
 
     chain = ChainStart[int]().flat_map(add_one)
@@ -100,11 +99,11 @@ def add_one(input_iterator: Iterator[int]) -> Iterator[int]:
 
 def test_chain_multiple_operations() -> None:
     def add_one(input_iterator: Iterator[int]) -> Iterator[int]:
-        while (element := take_next(input_iterator)) is not None:
+        for element in input_iterator:
             yield element + 1
 
     def multiply_two(input_iterator: Iterator[int]) -> Iterator[int]:
-        while (element := take_next(input_iterator)) is not None:
+        for element in input_iterator:
             yield element * 2
 
     def is_even(value: int) -> bool:

tests/core/test_chain.py

@@ -1,7 +1,7 @@
+from itertools import islice
 from typing import Iterable, Iterator, List
 
 from pipedata.core import ChainStart, StreamStart
-from pipedata.core.itertools import take_next, take_up_to_n
 
 
 def test_stream_to_list() -> None:
@@ -81,20 +81,19 @@ def is_even(value: int) -> bool:
 
 def test_stream_flat_map_identity() -> None:
     def identity(input_iterator: Iterator[int]) -> Iterator[int]:
-        while (element := take_next(input_iterator)) is not None:
-            yield element
+        yield from input_iterator
 
     result = StreamStart([0, 1, 2, 3]).flat_map(identity).to_list()
     assert result == [0, 1, 2, 3]
 
 
 def test_stream_flat_map_chain() -> None:
     def add_one(input_iterator: Iterator[int]) -> Iterator[int]:
-        while (element := take_next(input_iterator)) is not None:
+        for element in input_iterator:
             yield element + 1
 
     def multiply_two(input_iterator: Iterator[int]) -> Iterator[int]:
-        while (element := take_next(input_iterator)) is not None:
+        for element in input_iterator:
             yield element * 2
 
     result = (
@@ -105,7 +104,7 @@ def multiply_two(input_iterator: Iterator[int]) -> Iterator[int]:
 
 def test_stream_flat_map_growing() -> None:
     def add_element(input_iterator: Iterator[int]) -> Iterator[int]:
-        while (element := take_next(input_iterator)) is not None:
+        for element in input_iterator:
             yield element
             yield element + 1
 
@@ -115,7 +114,7 @@ def add_element(input_iterator: Iterator[int]) -> Iterator[int]:
 
 def test_stream_flat_map_shrinking() -> None:
     def add_two_values(input_iterator: Iterator[int]) -> Iterator[int]:
-        while batch := take_up_to_n(input_iterator, 2):
+        while batch := tuple(islice(input_iterator, 2)):
             yield sum(batch)
 
     result = StreamStart([0, 1, 2, 3, 4]).flat_map(add_two_values).to_list()