Flatten an array of integers in Python

Posted on

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])

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.

You can do this in two ways either via LBYL or EAFP.

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)
``````