flatten

import itertools
import unittest
from collections.abc import Iterable


def flatten(items):
    if isinstance(items, Iterable):
        yield from itertools.chain.from_iterable(flatten(item) for item in items)
    else:
        yield items


class TestFlatten(unittest.TestCase):
    def test_non_iterable(self):
        # expect return a iterable contain 5
        self.assertListEqual(list(flatten(5)), [5])

    def test_one_item(self):
        self.assertListEqual(list(flatten([123])), [123])

    def test_1d(self):
        self.assertListEqual(list(flatten([2, 3, 5, 7, 11])), [2, 3, 5, 7, 11])

    def test_2d(self):
        self.assertListEqual(list(flatten([[2, 3], [5, 7, 11]])), [2, 3, 5, 7, 11])

    def test_3d(self):
        self.assertListEqual(
            list(flatten([[[2, 2], [3, 3]], [[5, 5], [7, 7]], [[11, 13, 17]]])),
            [2, 2, 3, 3, 5, 5, 7, 7, 11, 13, 17],
        )

    def test_mixed_dimension(self):
        self.assertListEqual(
            list(flatten([2, [3, 5], [7, [11, 13]]])), [2, 3, 5, 7, 11, 13]
        )


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