flattener.py

import collections


def flatten(thing):
    """
    Given an arbitrarily nested list, flattens it into a giant list.

    Args:
        thing (:obj:`list`):

    Returns:
        Generator of elements found within the original :obj:`list` as a stream of elements.
    """

    for foo in thing:
        if isinstance(foo, collections.Iterable) and not isinstance(foo, basestring):
            for bar in flatten(foo):
                yield bar
        else:
            yield foo

test_flattener.py

import unittest
from flattener import flatten


class FlattenTest(unittest.TestCase):
    """ Test cases for flattener.flatten method """

    def testEmptyList(self):
        """ Test case: Flatten an empty list. Should return empty list """
        foo = []
        assert list(flatten(foo)) == foo, 'flatten an empty list did not return an empty list.'

    def testSimpleList(self):
        """ Test case: flatten an already flattened list. Should return original list """
        foo = [1, 2, 3, 4]
        assert list(flatten(foo)) == foo, 'flatten a shallow list did not return same shallow list.'

    def testComplexList(self):
        """ Test case: flatten a complex list, of varying types """
        foo = [[1, 2, [3]], 4]
        assert list(flatten(foo)) == [1, 2, 3, 4], 'flatten complex list from interview question did not return [1, 2, 3, 4].'

    def testDifferingTypesList(self):
        """
        Test case: flatten a list of varying types.
        Note, this case showed a critical bug when not checking for string types, which are iterables.
        """
        foo = [1, [2.0, 'three', ['four', ('five', 'six')]]]
        assert list(flatten(foo)) == [1, 2.0, 'three', 'four', 'five', 'six'], 'flatten mixed type list did not return flattened list.'

if __name__ == "__main__":
    unittest.main()  # run all tests