nested_arrays.py

TESTS = [
([1,[2,3],6],[1,2,3,6]),
([2,[[3],1,9]],[2,3,1,9]),
([32],[32]),
([[[[5]]]],[5]),
([[[[[[-1]]]]]],[-1]),
([1,[[[[[[[9]]]]]]],[[[[[[9]]]]]]],[1,9,9]),
([],[])
]

def flatten(nested):
	flat = []
	for i, val in enumerate(nested):
		if type(val) is list:
			flat.extend(flatten(val))
		else:
			flat.append(val)
	return flat

def test_flatten(nested,flat):
	flat.sort()
	test_output = flatten(nested)
	test_output.sort()

	if len(flat) != len (test_output):
		return False

	for i in range(len(flat)):
		if flat[i] != test_output[i]:
			return False
	return True

if __name__ == "__main__":
	err = False
	for i,tup in enumerate(TESTS):
		if not test_flatten(tup[0],tup[1]):
			err = True
			print "Error in case" + str(tup)
			break
	if not err:
		print "All tests passed!"