Problem
The goal is to flatten an array of nested arrays of integers.
For example [1, [2], [3, [4]]]
should return [1, 2, 3, 4]
I’m particularly looking for some feedback on the following:
- Is usage of
TypeError
exception justified? - Did I miss some valid edge cases in my tests?
Here is my solution including unit tests:
#!/usr/bin/env python
import unittest
def flatten_list(array):
""" Returns a flattened array of integers from arbitrarily nested arrays of integers.
>>> flatten_list([[1, 2, [3]], 4])
[1, 2, 3, 4]
"""
flat_array = []
for element in array:
if isinstance(element, int):
flat_array.append(element)
elif isinstance(element, list):
flat_array += flatten_list(element)
else:
raise TypeError("Unsupported type ({})".format(
type(element).__name__))
return flat_array
class FlattenListTestCase(unittest.TestCase):
"""Tests for flatten_list"""
def test_empty_list(self):
self.assertEqual(flatten_list([]), [])
def test_single_integer(self):
self.assertEqual(flatten_list([3]), [3])
def test_already_flattened(self):
self.assertEqual(flatten_list([1, 2, 3, 4]), [1, 2, 3, 4])
def test_single_element_multiple_nested(self):
self.assertEqual(flatten_list([[1]]), [1])
def test_arbitary(self):
self.assertEqual(flatten_list(
[1, [2], [3, 4], [[5]], 6]), [1, 2, 3, 4, 5, 6])
def test_empty_sublists(self):
self.assertEqual(flatten_list([1, [[], 2]]), [1, 2])
self.assertEqual(flatten_list([[]]), [])
def test_invalid_input(self):
with self.assertRaises(TypeError):
flatten_list(3)
with self.assertRaises(TypeError):
flatten_list([3, "2"])
with self.assertRaises(TypeError):
flatten_list(None)
if __name__ == '__main__':
unittest.main()
Solution
- Python normally promotes duck typing, yours is more statically typed.
- You may want to use
yield
, as your current code possibly uses more memory than needed. -
You can check if an item is iterable, and iterate it. Otherwise return the item.
Just using yield
can get:
def flatten_list(array):
for element in array:
if isinstance(element, int):
yield element
elif isinstance(element, list):
yield from flatten_list(element)
else:
raise TypeError("Unsupported type ({})".format(
type(element).__name__))
Using LBYL and duck typing can then get:
import collections
def flatten_list(array):
for element in array:
if isinstance(element, collections.abc.Iterable):
yield from flatten_list(element)
else:
yield element
Using EAFP and duck typing can instead get:
def flatten_list(array):
for element in array:
try:
element = iter(element)
except TypeError:
yield element
else:
yield from flatten_list(element)
Finally as a side note, if you expect the initial array to have a depth exceeding one thousand, then you may want to instead perform the stack manipulation yourself. So you don’t get a stack overflow error.
This can be done with something like:
def flatten(it):
to_yield = [iter(it)]
sentinel = object()
while to_yield:
item = next(to_yield[-1], sentinel)
if item is sentinel:
del to_yield[-1]
continue
try:
item = iter(item)
except TypeError:
yield item
else:
to_yield.append(item)