Sample code for Citrusbyte application

flatten.py

# flatten.py


def flatten(int_list):
    """
    Recursive function that transforms an arbitrarily nested list of integers into a flat list of integers.

    :param int_list: nested list of integers
    :return: list of integers
    """
    result = []

    # First, check to make sure `int_list` is a list. If not, raise an exception.
    if not isinstance(int_list, list):
        raise TypeError("{} is not a list".format(int_list))

    # Iterate over every item.  If it's an integer, add it to the resulting list.
    # If it's another list, recurse.  If it's neither, raise an exception.
    for item in int_list:
        if isinstance(item, list):
            result += flatten(item)
        elif isinstance(item, int):
            result.append(item)
        else:
            raise TypeError("{} is not an integer or list of integers.".format(item))

    return result

test_flatten.py

# test_flatten.py
# Provides unit tests for `flatten` module
#
# Execute with:
#   $ python -m unittest test_flatten.FlattenTestCase

from unittest import TestCase
from flatten import flatten


class FlattenTestCase(TestCase):

    def test_flatten(self):
        int_list = [[1, 2, [3]], 4]
        result = flatten(int_list)
        self.assertEqual(result, [1, 2, 3, 4])

        int_list = []
        result = flatten(int_list)
        self.assertEqual(result, [])

        int_list = [1, [], [[]], 2]
        result = flatten(int_list)
        self.assertEqual(result, [1, 2])

    def test_flatten_invalid_data(self):
        with self.assertRaises(TypeError):
            flatten(["1", "2", ["3"]])

        with self.assertRaises(TypeError):
            flatten(1)

        with self.assertRaises(TypeError):
            flatten("1")