pythonでflatten

Rubyにはネストされた配列を1次元配列にフラット化するflattenというメソッドがあるのですが、pythonでも使いたくなることがあるので、どんな方法があるか調べてみました。

やりたいこと

>>> 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]

Qiitaで見つけた方法

ジェネレータ関数にする

Qiitaで見つけた方法を参考にしてジェネレータ関数にしてみました。
リスト化するには list(flatten([[1, 2], [3, [4, 5]]])) のようにする必要があります。

def flatten(data):
    for item in data:
        if isinstance(item, list):
            yield from flatten(item)
        else:
            yield item

ジェネレータ内包表記にする

def flatten(data):
    return (element
            for item in data
            for element in (flatten(item) if isinstance(item, list) else [item]))

リスト内包表記にする

def flatten(data):
    return [element
            for item in data
            for element in (flatten(item) if isinstance(item, list) else [item])]

既存ライブラリにある関数を使う

ここで見つけました。
http://stackoverflow.com/questions/2158395/flatten-an-irregular-list-of-lists-in-python

from compiler.ast import flatten

sum関数を使ってリストのリストを深さ1だけflattenする

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

深さ指定オプション付き関数

Rubyではフラット化する深さを指定できるので、対応してみました。

flatten.py

#!/usr/bin/env python3

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 isinstance(item, list)
                            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.