diff --git a/workspaces/adventure-pack/goodies/python3/src/flatten/__init__.py b/workspaces/adventure-pack/goodies/python3/src/flatten/__init__.py index 8273e6fc..c541eb1d 100644 --- a/workspaces/adventure-pack/goodies/python3/src/flatten/__init__.py +++ b/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"] -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 diff --git a/workspaces/adventure-pack/goodies/python3/src/flatten/test.py b/workspaces/adventure-pack/goodies/python3/src/flatten/test.py index 281e067d..b09614fb 100644 --- a/workspaces/adventure-pack/goodies/python3/src/flatten/test.py +++ b/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] diff --git a/workspaces/adventure-pack/src/app/__tests__/__snapshots__/equip-test.ts.snap b/workspaces/adventure-pack/src/app/__tests__/__snapshots__/equip-test.ts.snap index 316c6ea5..f7a3641e 100644 --- a/workspaces/adventure-pack/src/app/__tests__/__snapshots__/equip-test.ts.snap +++ b/workspaces/adventure-pack/src/app/__tests__/__snapshots__/equip-test.ts.snap @@ -2835,17 +2835,26 @@ exports[`App can equip single goody: Python 3 flatten 1`] = ` # Running at: https://example.com/ from typing import Generator, TypeVar +from math import inf T = TypeVar("T") NestedList = T | list["NestedList"] -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 ########################### END ADVENTURE PACK CODE ############################" `; diff --git a/workspaces/adventure-pack/src/app/__tests__/__snapshots__/render-test.ts.snap b/workspaces/adventure-pack/src/app/__tests__/__snapshots__/render-test.ts.snap index b96fdd53..f0a6e8f0 100644 --- a/workspaces/adventure-pack/src/app/__tests__/__snapshots__/render-test.ts.snap +++ b/workspaces/adventure-pack/src/app/__tests__/__snapshots__/render-test.ts.snap @@ -1579,17 +1579,26 @@ exports[`App can render goody: Python 3 UnionFind 1`] = ` exports[`App can render goody: Python 3 flatten 1`] = ` "from typing import Generator, TypeVar +from math import inf T = TypeVar("T") NestedList = T | list["NestedList"] -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" `; exports[`App can render goody: Python 3 int.digits 1`] = `