def test_single_nest_list(self):
l = [1, 2, 3, [4]]
flat = list(_.flatten(l))
self.assertEqual([1, 2, 3, 4], flat)
l = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
flat = list(_.flatten(l))
self.assertEqual(flat, [1, 2, 3, 4, 5, 6, 7, 8, 9])
def test_single_nest_generator(self):
l = [range(0, 3), range(3, 6), range(6, 9)]
flat = list(_.flatten(l))
self.assertEqual(flat, [0, 1, 2, 3, 4, 5, 6, 7, 8])
l = chain(range(0, 3), range(3, 6), range(6, 9))
flat = list(_.flatten(l))
self.assertEqual(flat, [0, 1, 2, 3, 4, 5, 6, 7, 8])
def test_deep_nested_generator(self):
gen = range(0, 3)
squared_gen = lambda: map(lambda x: x * x, range(0, 3))
nested_gen = chain(gen, gen)
nested_gen2 = chain(squared_gen(), squared_gen())
deep_nested_generator = chain(nested_gen, nested_gen2)
flat = list(_.flatten(deep_nested_generator))
self.assertEqual(flat, [0, 1, 2, 0, 1, 2, 0, 1, 4, 0, 1, 4])
print(flat)
def test_flat_list(self):
l = [1, 2, 3, 4]
flat = list(_.flatten(l))
self.assertEqual(l, flat)
def test_deep_nested_list(self):
l = [[1, 2], [3, 4], [[5, 6], 7, 8]]
flat = list(_.flatten(l))
self.assertEqual(flat, [1, 2, 3, 4, 5, 6, 7, 8])
def test_flat_generator(self):
l = range(0, 10)
flat = _.flatten(range(0, 10))
self.assertEqual(list(l), list(flat))