[python] filterで生成したオブジェクトを再利用してはならない

 pythonのbuilt-in関数であるfilterは大変便利なものですが、とんでもない落とし穴があります。

落とし穴

 まず、有効なレコードをフィルターし、その中から更にいくつかの属性に分類する、という仮想コードを書いてみましょう。

records = [
  {
    'enabled': False,
    'attribute': 'a',
  },
  {
    'enabled': True,
    'attribute': 'a',
  },
  {
    'enabled': True,
    'attribute': 'b',
  },
]

enabled_records = filter(lambda record: record['enabled'], records)
attribute_a_records = filter(lambda record: record['attribute'] == 'a', enabled_records)
attribute_b_records = filter(lambda record: record['attribute'] == 'b', enabled_records)

print(list(attribute_a_records))  # -> [{'enabled': True, 'attribute': 'a'}]
print(list(attribute_b_records))  # -> []

 attribute_b_recordsの値が変ですね。想定では、 [{'enabled': True, 'attribute': 'b'}] となるはずでした。

filterの正体

 filterの正体は、公式ドキュメントを読むとわかるとおり、generatorです。これは、listやfor〜in文で実行されたときに初めて評価(遅延評価)されることを意味しています。

generatorとは?

 細かい説明は割愛しますが、以下のコードをご覧ください。これはフィボナッチ数列を引数分生成するgeneratorです。要素が取得される毎にyieldで実行が停止する、というイメージが近いと思います。

def fibonacci(n):
  a = 0
  b = 1
  for i in range(n):
    c = a + b
    a = b
    b = c
    yield c

 実際に使ってみます。

print(list(fibonacci(10)))  # -> [1, 2, 3, 5, 8, 13, 21, 34, 55, 89]

 では、一度generatorを実行し、変数に格納してから、それを再利用してみましょう。

f = fibonacci(10)
print(list(f))  # -> [1, 2, 3, 5, 8, 13, 21, 34, 55, 89]
print(list(f))  # -> []

 ここで、先程のgenerator関数を思い出してみましょう。yieldで一旦実行が停止する、と考えると理解がスムーズです。

 fibonacci関数のfor文で設定されたrange(n)、この場合はn=10ですが、これを使い果たしてfor文が終了したため、次に返すべきyieldの行が見つからず、要素を取得することができなかったというわけです。

 listなどと大きく違う点は、再度要素を先頭から取得し直すには、基本的にはgeneratorを作り直すしかないということです。また、generatorはその名の通り、生成器であるため、事前に要素の数を知るなどといった、遠い未来を予測することができません。

filterの実装

 実は、公式ドキュメントに実装が載っています。非常に簡潔ですから、ご紹介しておきます。

なお、filter(function, iterable) は、関数が None でなければジェネレータ式 (item for item in iterable if function(item)) と同等で、関数が None なら (item for item in iterable if item) と同等です。

 わかりやすいかどうかはともかく、ワンライナーでない分解したコードを書いてみましょう。

def _filter(fn, iter):
  for item in iter:
    if fn is not None:
      if fn(item):
        yield item
    else:
      if item is not None:
        yield item

 これがfilterの実装です。実際にテストしてみます。

enabled_records = _filter(lambda record: record['enabled'], records)
attribute_a_records = _filter(lambda record: record['attribute'] == 'a', enabled_records)
attribute_b_records = _filter(lambda record: record['attribute'] == 'b', enabled_records)

print(list(attribute_a_records))  # -> [{'enabled': True, 'attribute': 'a'}]
print(list(attribute_b_records))  # -> []

 先程と同様の結果になりました。

filterを使うときに気をつけること

 毎回、必ずlistなどの実体からfilterすることを心がけてください。そうでなければ、思わぬ(そして、デバッグ困難な)結果を生むことになります。

def create_enabled_filter():
  return filter(lambda record: record['enabled'], records)

attribute_a_records = filter(lambda record: record['attribute'] == 'a', create_enabled_filter())
attribute_b_records = filter(lambda record: record['attribute'] == 'b', create_enabled_filter())

print(list(attribute_a_records))  # -> [{'enabled': True, 'attribute': 'a'}]
print(list(attribute_b_records))  # -> [{'enabled': True, 'attribute': 'b'}]

 毎回、新しいfilterオブジェクトを作った結果、想定通りの振る舞いが確認できました。

 もしくは、途中経過はlistにしてしまいましょう。こちらのほうが簡潔かもしれません。

enabled_records = list(filter(lambda record: record['enabled'], records))
attribute_a_records = filter(lambda record: record['attribute'] == 'a',enabled_records)
attribute_b_records = filter(lambda record: record['attribute'] == 'b',enabled_records)

print(list(attribute_a_records))  # -> [{'enabled': True, 'attribute': 'a'}]
print(list(attribute_b_records))  # -> [{'enabled': True, 'attribute': 'b'}]

まとめ

 filterの危険な使い方について解説しましたが、実はitertoolsにもgeneratorが盛りだくさんです。built-in関数は大変便利なものですが、必要に応じてドキュメントを読み、正しく理解して使うことをおすすめします。