find もどき

ディレクトリをたどってファイルを処理するとき毎回書いてしまう os.walk 。もう自然と手が for root, dirs, files in os.walk(dir): って動くようになってしまった。でもよくよく考えたら os.walk ジェネレーターって汎用性が高い分、使いにくい。たとえば拡張子 .py のファイルを探すのに os.walk を使ってしまうと 2 重 for 文に if 文、そして fnmatch.fnmatch か os.path.splitext もしくは str.endswith がいる。

find コマンドを使えば find dir -name '*.py' でできてしまうことなのに。とはいえ subprocess.Popen(["find", dir, "-name", "*.py"]) する気はあまりおこらない。

というわけで書いてみた。find(dir, u'*.py') とか find(dir, re.compile(ur'.*\.py$')) で拡張子 .py のファイルを処理できる。

import fnmatch
import functools
import os
import re

def find(path, pattern):
    if isinstance(pattern, basestring):
        match = functools.partial(fnmatch.fnmatch, pat=pattern)
    elif isinstance(pattern, re._pattern_type):
        match = pattern.match
    elif callable(pattern):
        match = pattern
    else:
        raise TypeError

    for root, dirs, files in os.walk(path):
        for file_ in files:
            if match(file_):
                yield os.path.join(root, file_)

第 1 引数 path は探るディレクトリのパス、第 2 引数 pattern は探すファイル名のパターンで文字列か正規表現オブジェクトか呼び出し可能オブジェクト。文字列なら Unix シェルのワイルドカード使用っぽいマッチ処理を行う。正規表現オブジェクトならその match メソッドを使ったマッチ処理をする（ match で作ったけれど場合によっては search のほうが対応幅が広くて良いかもしれない）。呼び出し可能オブジェクトの場合は、引数としてファイル名１つを受け取るものであることを期待。

今後 os.walk はこまかい探索するときだけご登場願う予定。