flatten in python

Ruby has a method called flatten that flattens a nested array into a one-dimensional array, but I sometimes want to use it in python, so I tried to find out what kind of method there is.

Thing you want to do

>>> flatten([[1, 2], [3, [4, 5]]])
[1, 2, 3, 4, 5]
>>> flatten([1, [2, 3], [[4, [5, 6]], 7]])
[1, 2, 3, 4, 5, 6, 7]

How to find it on Qiita

http://qiita.com/kento1218@github/items/f3baf574aadb3d1cbeae

Make it a generator function

I tried to make it a generator function by referring to the method found in Qiita. To list it, you need to do something like list (flatten ([[1, 2], [3, [4, 5]]])).

def flatten(data):
    for item in data:
        if hasattr(item, '__iter__'):
            for element in flatten(item):
                yield element
        else:
            yield item

Generate generator comprehension

def flatten(data):
    return (element
            for item in data
            for element in (flatten(item) if hasattr(item, '__iter__') else [item]))

List comprehension

def flatten(data):
    return [element
            for item in data
            for element in (flatten(item) if hasattr(item, '__iter__') else [item])]

Use functions in existing libraries

I found it here. http://stackoverflow.com/questions/2158395/flatten-an-irregular-list-of-lists-in-python

from compiler.ast import flatten

Flatten the list of lists by depth 1 using the sum function

>>> data = [[1, 2], [3], [4, 5, [6]]]
>>> sum(data, [])
[1, 2, 3, 4, 5, [6]]

Function with depth specification option

In Ruby, you can specify the depth to flatten, so I tried to support it.

flatten.py

#!/usr/bin/env python
# -*- coding:utf-8 -*-

def flatten(data, depth=-1):
    """
    flatten(data) -> list
    flatten(data, depth) -> list

    Return flatted data of list or tupple as list.

    >>> data = [[1, 2], [3, [4, 5, [6]]]]
    >>> flatten(data)
    [1, 2, 3, 4, 5, 6]
    >>> flatten(data, 0)
    [[1, 2], [3, [4, 5, [6]]]]
    >>> flatten(data, 1)
    [1, 2, 3, [4, 5, [6]]]
    >>> flatten(data, 2)
    [1, 2, 3, 4, 5, [6]]
    >>> flatten(data, 3)
    [1, 2, 3, 4, 5, 6]
    """
    return [element
            for item in data
            for element in (flatten(item, depth - 1)
                            if depth != 0 and hasattr(item, '__iter__')
                            else [item])
            ]

$ python -m doctest -v flatten.py
Trying:
    data = [[1, 2], [3, [4, 5, [6]]]]
Expecting nothing
ok
Trying:
    flatten(data)
Expecting:
    [1, 2, 3, 4, 5, 6]
ok
Trying:
    flatten(data,0)
Expecting:
    [[1, 2], [3, [4, 5, [6]]]]
ok
Trying:
    flatten(data, 1)
Expecting:
    [1, 2, 3, [4, 5, [6]]]
ok
Trying:
    flatten(data, 2)
Expecting:
    [1, 2, 3, 4, 5, [6]]
ok
Trying:
    flatten(data, 3)
Expecting:
    [1, 2, 3, 4, 5, 6]
ok
1 items had no tests:
    flatten
1 items passed all tests:
   6 tests in flatten.flatten
6 tests in 2 items.
6 passed and 0 failed.
Test passed.