A Function That Flattens a Nested Array

Asked for a function that flattens a nested array in Python. Here's what came out after a bit of enterprise seasoning.

from typing import Any, Iterable, List, Union, Optional
from abc import ABC, abstractmethod

# Type alias for a nested list of arbitrary depth
NestedList = Union[Any, List['NestedList']]


class FlattenStrategy(ABC):
    """Abstract base class for all flattening strategies."""

    @abstractmethod
    def flatten(self, data: Iterable[Any]) -> List[Any]:
        pass


class RecursiveFlattenStrategy(FlattenStrategy):
    """Flattens a nested iterable using recursion."""

    def __init__(self, max_depth: Optional[int] = None):
        # Store the maximum recursion depth (None means unlimited)
        self.max_depth = max_depth

    def flatten(self, data: Iterable[Any]) -> List[Any]:
        # Delegate to the internal recursive helper
        return self._flatten_helper(data, current_depth=0)

    def _flatten_helper(self, data: Iterable[Any], current_depth: int) -> List[Any]:
        result: List[Any] = []

        if data is None:
            # Guard against None input
            return result

        for item in data:
            if isinstance(item, list) and (self.max_depth is None or current_depth < self.max_depth):
                # Recurse into nested lists
                result.extend(self._flatten_helper(item, current_depth + 1))
            else:
                # Append non-list items directly
                result.append(item)

        return result


def flatten(nested: NestedList, max_depth: Optional[int] = None) -> List[Any]:
    """Public API for flattening a nested list."""
    if not isinstance(nested, list):
        raise TypeError(f"Expected a list, got {type(nested).__name__}")

    strategy = RecursiveFlattenStrategy(max_depth=max_depth)
    try:
        return strategy.flatten(nested)
    except RecursionError:
        # Fall back to iterative flattening if recursion limit is hit
        return list(Iterable.flat_map(nested))


if __name__ == "__main__":
    example = [1, [2, [3, [4, 5]], 6], 7, [8, [9]]]
    print(flatten(example))

Code Review

1. Line 5. NestedList = Union[Any, List['NestedList']] simplifies to just Any, since Any already includes lists. This type alias contributes nothing but vibes.

2. Lines 8-13. An ABC with a single concrete implementation. We invented the Strategy pattern for a six line function. Congrats.

3. Lines 30-32. Guarding against None when the public entry point already raises TypeError for non-lists. This branch is unreachable.

4. Lines 46-48. The type check rejects tuples, generators, and anything else iterable, which contradicts the Iterable annotations everywhere else. Pick a lane.

5. Lines 50-54. Iterable.flat_map does not exist in Python. This fallback would crash harder than the RecursionError it claims to handle.

6. Lines 19-21. Comment restates that self.max_depth stores max_depth. Riveting documentation.

7. Line 4. 'Type alias for a nested list of arbitrary depth' next to a definition that boils down to Any. The comment is doing more work than the code.

8. Lines 57-59. After 55 lines of ceremony, the actual usage is one line. The ratio of scaffolding to output speaks for itself.