From 7c5674e121226f83005d9c03fe79aa02a8dc9c3b Mon Sep 17 00:00:00 2001
From: Martin Bernstorff <martinbernstorff@gmail.com>
Date: Thu, 7 Mar 2024 20:48:43 +0100
Subject: [PATCH] fix(#137): implement lazy flattening

Fixes #137
---
 iterpy/iter.py | 30 +++++++++++++++++++-----------
 1 file changed, 19 insertions(+), 11 deletions(-)

diff --git a/iterpy/iter.py b/iterpy/iter.py
index 6212d00..cc7da37 100644
--- a/iterpy/iter.py
+++ b/iterpy/iter.py
@@ -3,10 +3,12 @@
 import copy
 import multiprocessing
 from collections import defaultdict
-from collections.abc import Callable, Iterable, Iterator, Sequence
 from functools import reduce
 from itertools import islice
-from typing import Generic, TypeVar
+from typing import TYPE_CHECKING, Any, Generator, Generic, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable, Iterator
 
 T = TypeVar("T")
 S = TypeVar("S")
@@ -97,16 +99,22 @@ def groupby(self, func: Callable[[T], str]) -> Iter[tuple[str, list[T]]]:
         return Iter(tuples)
 
     def flatten(self) -> Iter[T]:
-        values: list[T] = []
-        for i in self._iterator:
-            if isinstance(i, Sequence) and not isinstance(i, str):
-                values.extend(i)  # type: ignore
-            elif isinstance(i, Iter):
-                values.extend(i.to_list())  # type: ignore
+        depth = 1
+
+        def walk(node: Any, level: int) -> Generator[T, None, None]:
+            if (level > depth) or isinstance(node, str):
+                yield node  # type: ignore
+                return
+            try:
+                tree = iter(node)
+            except TypeError:
+                yield node
+                return
             else:
-                values.append(i)
+                for child in tree:
+                    yield from walk(child, level + 1)
 
-        return Iter(values)
+        return Iter(walk(self, level=0))  # type: ignore
 
     def take(self, n: int = 1) -> Iter[T]:
         return Iter(islice(self._iter, n))
@@ -141,7 +149,7 @@ def find(self, func: Callable[[T], bool]) -> T | None:
                 return value
         return None
 
-    def copy(self) -> Iter[T]:
+    def clone(self) -> Iter[T]:
         return copy.deepcopy(self)
 
     def zip(self, other: Iter[S]) -> Iter[tuple[T, S]]: