fix: Iter[iter].flatten() should return Iter (#94)

fix: Iter[iter].flatten() should return Iter update _iter.py and test_iter.py
MartinBernstorff · Jan 23, 2024 · c8790a7 · c8790a7
2 parents c2ea2d3 + e7f776c
commit c8790a7
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/iterpy/_iter.py b/iterpy/_iter.py
@@ -114,6 +114,8 @@ def flatten(self) -> "Iter[T]":
         for i in self._iterator:
             if isinstance(i, Sequence) and not isinstance(i, str):
                 values.extend(i)
+            elif isinstance(i, Iter):
+                values.extend(i.to_list())
             else:
                 values.append(i)
 

diff --git a/iterpy/test_iter.py b/iterpy/test_iter.py
@@ -124,7 +124,8 @@ def test_flatten_iterator(self):
     def test_flatten_iter_iter(self):
         iterator: Iter[int] = Iter([1, 2])
         nested_iter: Iter[Iter[int]] = Iter([iterator])
-        unnested_iter: Iter[int] = nested_iter.flatten()  # noqa: F841, RUF100
+        unnested_iter: Iter[int] = nested_iter.flatten()
+        assert unnested_iter.to_list() == [1, 2]
 
     def test_flatten_str(self):
         test_input: list[str] = ["abcd"]