Flatten an array of integers in Python

Question

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

Answer 1

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