flatten-plus-test.py

"""
flatten.py

Problem:
given an arbitarily nested list return a flattened list

Unit test included in the module itself for simplicity

Works with python2 and python3
"""
import unittest


def flatten(nested):
    "returns a flattened version of the provided list"
    if nested is None:
        raise TypeError("Cannot flatten a None object")

    if not isinstance(nested, list):
        raise TypeError("Cannot flatten a non list object")

    flattened = []
    for element in nested:
        if isinstance(element, list):
            flattened.extend(flatten(element))
        else:
            flattened.append(element)
    return flattened


class TestFlattenFunction(unittest.TestCase):

    def test_none_parameter(self):
        with self.assertRaises(TypeError):
            flatten(None)

    def test_not_list_parameter(self):
        with self.assertRaises(TypeError):
            flatten('not a list')

    def test_flatten(self):
        to_flatten = [1, 2, 3, [4], 5, [6, [7, 8, [9, [10]]]]]
        self.assertEqual(flatten(to_flatten), [1, 2, 3, 4, 5, 6, 7, 8, 9, 10])


if __name__ == '__main__':
    unittest.main()