[Python] リスト内包表記 大解剖

ネット上にリスト内包表記を解説する記事はいくつかありますが、簡単にヒットする記事ではざっくりとした内容でしか紹介されないケースが多いので、この記事ではPythonのAST(抽象構文木)を読み解くことでリスト内包表記でどんな書き方ができるか解説しようと思います。

※どんな書き方ができるかを知りたいだけの方は最初の方は読み飛ばしてください。

この記事はINIAD(東洋大学 情報連携学部) Advent Calendar 2020の記事です。

環境

Python 3.9.0

リスト内包表記とは

Pythonistaの皆様はご存知かもしれませんが、普段Pythonを利用しない方へ向けて軽く説明します。

以下は通常のforループでリストに一つずつ要素を追加する例です。

l = []
for i in range(9):
  l.append(i)

これをリスト内包表記を用いて書き換えると以下のようになります。

l = [i for i in range(9)]

また、ifで条件を付け加えることで要素のフィルタリングをすることもできます。

l = [i for i in range(9) if i % 2 == 0]

この場合、2で割り切れる要素のみがリストに含まれます。

リスト内包表記は、forループを利用し要素を追加するよりも スマート かつ高速に動作させることができる記法です。

というのも、Pythonインタプリタでは、リスト内包表記による要素の追加にはLIST_APPENDという専用のバイトコードが割り当てられており、forループで都度appendメソッドを呼び出す関数呼び出しのコストやリストのappend属性を探すコストを無視できるので、その分高速に動作するという仕組みです。

リスト内包表記のASTを得る

ASTを得るためにastモジュールを利用してリスト内包表記([i for i in range(9)])のASTを得てみたいと思います。

import ast

source_code = """
[i for i in range(9)]
"""

source_tree = ast.parse(source_code)

print(ast.dump(source_tree))
Module(
    body=[
        Expr(
            value=ListComp(
                elt=Name(id='i', ctx=Load()),
                generators=[
                    comprehension(
                        target=Name(id='i', ctx=Store()),
                        iter=Call(
                            func=Name(id='range', ctx=Load()),
                            args=[Constant(value=9, kind=None)],
                            keywords=[]
                            ),
                        ifs=[],
                        is_async=0
                    )
                ]
            )
        )
    ],
    type_ignores=[]
)

※見やすくするために改行を入れて整形しています。

ASTを読み解く

ASTの中身を見ると変数のirange関数の呼び出しが含まれていることから、リスト内包表記はListCompクラスで表されていることがわかります。

ListCompeltが[i for i in range]の部分、

ListCompgeneratorsにリストでcomprehensionクラスのインスタンスが保存されていることがわかります。
comprehensionクラスには[i for i in range(9)]の部分が保存されていることが見て取れます。
comprehensionクラスはListCompgeneratorsにリストで保存されていることからfor i in range(9)部分を繰り返し書くことが可能であることが推測できます。

comprehensionクラスのifsにリストで何かが保存されそうです。
名前から条件ifであるような気がします。

という事で確認してみます。

[i for i in range(9) if i % 2 == 0]

のASTは以下のようになりました。

Module(
    body=[
        Expr(
            value=ListComp(
                elt=Name(id='i', ctx=Load()),
                generators=[
                    comprehension(
                        target=Name(id='i', ctx=Store()),
                        iter=Call(
                            func=Name(id='range', ctx=Load()),
                            args=[Constant(value=9, kind=None)],
                            keywords=[]
                            ),
                        ifs=[
                            Compare(
                                left=BinOp(
                                    left=Name(id='i', ctx=Load()), op=Mod(),
                                    right=Constant(value=2, kind=None)
                                    ),
                                ops=[Eq()],
                                comparators=[Constant(value=0, kind=None)]
                            )
                        ],
                        is_async=0
                    )
                ]
            )
        )
    ],
    type_ignores=[]
)

想定通りifsには条件ifが保存されていました。

わざわざリストで保存されているということは、複数の条件を設定することができることが推測できます。

ということでASTを見た限りfor i in range(9)if i % 2 == 2の部分は複数かけそうな雰囲気なのでそれぞれ実際に試してみます。

forを複数並べる(多重ループ)

気になる部分その1

複数forを並べてみる。

[i for i in range(3) for j in range(3)]
[0, 0, 0, 1, 1, 1, 2, 2, 2]

これだけだと判りづらいのでちょっと書き換えます。

[(i, j) for i in range(3) for j in "ABC"]
[(0, 'A'), (0, 'B'), (0, 'C'), (1, 'A'), (1, 'B'), (1, 'C'), (2, 'A'), (2, 'B'), (2, 'C')]

要素数から回ったループの回数は3*3の9回なことが伺えます。

また、次の二重ループのコードと同じ値が得られることから後ろに書いたループが内側に入ることがわかります。

l = []
for i in range(3):
  for j in "ABC":
     l.append((i, j))

print(l)

for ~~ 部分は3回でも4回でも繰り返せます。

[(i, j, k, l) for i in range(3) for j in "ABC" for k in "abc" for l in "XYZ"]
[
(0, 'A', 'a', 'X'), (0, 'A', 'a', 'Y'), (0, 'A', 'a', 'Z'), (0, 'A', 'b', 'X'), 
(0, 'A', 'b', 'Y'), (0, 'A', 'b', 'Z'), (0, 'A', 'c', 'X'), (0, 'A', 'c', 'Y'), 
(0, 'A', 'c', 'Z'), (0, 'B', 'a', 'X'), (0, 'B', 'a', 'Y'), (0, 'B', 'a', 'Z'), 
(0, 'B', 'b', 'X'), (0, 'B', 'b', 'Y'), (0, 'B', 'b', 'Z'), (0, 'B', 'c', 'X'), 
(0, 'B', 'c', 'Y'), (0, 'B', 'c', 'Z'), (0, 'C', 'a', 'X'), (0, 'C', 'a', 'Y'), 
(0, 'C', 'a', 'Z'), (0, 'C', 'b', 'X'), (0, 'C', 'b', 'Y'), (0, 'C', 'b', 'Z'), 
(0, 'C', 'c', 'X'), (0, 'C', 'c', 'Y'), (0, 'C', 'c', 'Z'), (1, 'A', 'a', 'X'), 
(1, 'A', 'a', 'Y'), (1, 'A', 'a', 'Z'), (1, 'A', 'b', 'X'), (1, 'A', 'b', 'Y'), 
(1, 'A', 'b', 'Z'), (1, 'A', 'c', 'X'), (1, 'A', 'c', 'Y'), (1, 'A', 'c', 'Z'), 
(1, 'B', 'a', 'X'), (1, 'B', 'a', 'Y'), (1, 'B', 'a', 'Z'), (1, 'B', 'b', 'X'), 
(1, 'B', 'b', 'Y'), (1, 'B', 'b', 'Z'), (1, 'B', 'c', 'X'), (1, 'B', 'c', 'Y'), 
(1, 'B', 'c', 'Z'), (1, 'C', 'a', 'X'), (1, 'C', 'a', 'Y'), (1, 'C', 'a', 'Z'), 
(1, 'C', 'b', 'X'), (1, 'C', 'b', 'Y'), (1, 'C', 'b', 'Z'), (1, 'C', 'c', 'X'), 
(1, 'C', 'c', 'Y'), (1, 'C', 'c', 'Z'), (2, 'A', 'a', 'X'), (2, 'A', 'a', 'Y'), 
(2, 'A', 'a', 'Z'), (2, 'A', 'b', 'X'), (2, 'A', 'b', 'Y'), (2, 'A', 'b', 'Z'), 
(2, 'A', 'c', 'X'), (2, 'A', 'c', 'Y'), (2, 'A', 'c', 'Z'), (2, 'B', 'a', 'X'), 
(2, 'B', 'a', 'Y'), (2, 'B', 'a', 'Z'), (2, 'B', 'b', 'X'), (2, 'B', 'b', 'Y'), 
(2, 'B', 'b', 'Z'), (2, 'B', 'c', 'X'), (2, 'B', 'c', 'Y'), (2, 'B', 'c', 'Z'), 
(2, 'C', 'a', 'X'), (2, 'C', 'a', 'Y'), (2, 'C', 'a', 'Z'), (2, 'C', 'b', 'X'), 
(2, 'C', 'b', 'Y'), (2, 'C', 'b', 'Z'), (2, 'C', 'c', 'X'), (2, 'C', 'c', 'Y'), 
(2, 'C', 'c', 'Z')
]

要素をフィルタリングする(filter関数的使い方)

冒頭でも紹介しましたが、リスト内包表記ではifを付け加えることで要素をフィルタリングできます。

[i for i in range(12) if i != 3]
[0, 1, 2, 4, 5, 6, 7, 8, 9, 10, 11]

ifを並べて複数の条件を指定する

というわけで気になる部分その2 ifを複数並べてみるパターンです。

[i for i in range(12) if i != 0 if i != 3]
[1, 2, 4, 5, 6, 7, 8, 9, 10, 11]

両方の条件に合致する要素のみが含まれるようになるみたいです。

ifを複数並べて書けるなんて面白いですね。
ちなみに以下のようにandでつないで書いても意味的には等価です。

[i for i in range(12) if i != 0 and i != 1]

速度面ではほんのわずかにですがifを複数並べて書いたほうが速いみたいです。
(andで繋ぐとPythonの式としてandが行われるが複数ifを並べる方はcのコードで評価されるので)

少しでも速度面が気になる方はifを複数並べて書いてみるといいかもしれません。

要素に関数を適用する(map関数的使い方)

ちょっと本題からは外れる内容ですが左側のiを関数に渡してあげると全ての要素にその関数を適用するmap関数的な使い方もできます。

def double(n):
  return n * 2

[double(i) for i in range(9)]

# [0, 2, 4, 6, 8, 10, 12, 14, 16]

関数以外にも三項演算もできるので、こんなこともできます。

[i if i % 2 == 0 else 0 for i in range(9)]
[0, 0, 2, 0, 4, 0, 6, 0, 8]

応用

多重ループでフィルタリングも行う

外側のループに条件を指定するパターン
[(i, j) for i in range(3) if i % 2 == 0 for j in "ABC"]
[(0, 'A'), (0, 'B'), (0, 'C'), (2, 'A'), (2, 'B'), (2, 'C')]
内側のループに条件を指定するパターン
[(i, j) for i in range(3) for j in "ABC" if j != "B"]
[(0, 'A'), (0, 'C'), (1, 'A'), (1, 'C'), (2, 'A'), (2, 'C')]

つまり[(i, j) for i in range(3) if i % 2 == 0 for j in "ABC" if j != "B"]がセット

最後に

  • for を複数並べることで多重ループが表現できる
  • if を複数並べることでand条件でフィルタリングできる

以上、リスト内包表記のASTからリスト内包表記で可能な記法を紹介してみました!