feat: add depth parameter to the Python flatten function

Author: skdthread

workspaces/adventure-pack/goodies/python3/src/flatten/__init__.py

@@ -1,12 +1,21 @@
 from typing import Generator, TypeVar
+from math import inf
 
 T = TypeVar("T")
-NestedList = T | list["NestedList"]
+NestedList = T | list["nested_list"]
 
 
-def flatten(self: NestedList) -> Generator[T, None, None]:
-    for i in self:
-        if isinstance(i, list):
-            yield from flatten(i)
+def flatten(
+    nested_list: NestedList, depth: int = inf
+) -> Generator[T, None, None]:
+    """Flatten a nested list up to the specified depth.
+
+    Args:
+        nested_list: The nested list to flatten
+        depth: The maximum depth to flatten. Default is infinity (flatten all levels).
+    """
+    for item in nested_list if isinstance(nested_list, list) else [nested_list]:
+        if isinstance(item, list) and depth > 0:
+            yield from flatten(item, depth - 1)
         else:
-            yield i
+            yield item

workspaces/adventure-pack/goodies/python3/src/flatten/test.py

@@ -14,3 +14,38 @@ def test_flatten_random_types():
     flattened = list(flatten(l))
     expected = ["hi", -80, "c", 0.0, (3, 3), dictionary, 6, 5, 4, 3, 2, 1, 0]
     assert flattened == expected
+
+
+def test_flatten_with_depth_0():
+    # Test with depth=0 (no flattening)
+    input_ = [1, [2, [3, [4, 5]]]]
+    assert list(flatten(input_, depth=0)) == input_
+
+
+def test_flatten_with_depth_2():
+    # Test with depth=2
+    input_ = [1, [2, [3, [4, 5]]]]
+    expected = [1, 2, 3, [4, 5]]
+    assert list(flatten(input_, depth=2)) == expected
+
+
+def test_flatten_with_depth_greater_than_needed():
+    # Test with depth larger than needed
+    input_ = [1, [2, [3, [4, 5]]]]
+    expected = [1, 2, 3, 4, 5]
+    assert list(flatten(input_, depth=10)) == expected
+
+
+def test_flatten_with_empty_lists():
+    # Test with empty lists at different levels
+    input_ = [[], [1, [2, [], [3, []]], 4], []]
+    assert list(flatten(input_, depth=1)) == [1, [2, [], [3, []]], 4]
+    assert list(flatten(input_, depth=2)) == [1, 2, [], [3, []], 4]
+    assert list(flatten(input_)) == [1, 2, 3, 4]
+
+
+def test_flatten_non_list_input():
+    # Test with non-list input
+    assert list(flatten(42)) == [42]
+    assert list(flatten("hello")) == ["hello"]
+    assert list(flatten(None)) == [None]