Flatten a nested array

flatten_nested.py

def flatten(arr):
  """ 
  Flatten a nested array

  >>> flatten([1, 2, 3])
  [1, 2, 3]

  >>> flatten([1,2,3,[1,2,3]])
  [1, 2, 3, 1, 2, 3]

  >>> flatten([1,2,3,[1,2,[4,5,6],3]])
  [1, 2, 3, 1, 2, 4, 5, 6, 3]
  """
  flattened_arr = []
  for elem in arr:
    if not isinstance(elem, list):
      flattened_arr.append(elem)
    else:
      [flattened_arr.append(i) for i in flatten(elem)]
  return flattened_arr

if __name__ == "__main__":
  import doctest
  doctest.testmod()

This problem could also be a good use case for the yield statement.

Neat! Was looking for this. Good to learn about doctest as well.