特定の拡張子のついたファイルを取得してみた改良案

元ネタは「特定の拡張子のついたファイルを取得してみた - 牌語備忘録 -pygo」。

たぶん、牌語備忘録さんがさがしていたのは fnmatch.filter だと思ったのでトラックバックを。複数のほうはこれといったものはないけれど、 glob じゃなくて os.listdir と fnmatch.fnmatch で作り直して、と。

このようなものができました。 os.chdir しなくてもよくなりました。

import os
import fnmatch

# 一つの拡張子のファイル名を取得したい場合
def hoge(dir_path, ext):
    return fnmatch.filter(os.listdir(dir_path), ext)

# 複数の拡張子のファイル名を取得したい場合
def ifuga(dir_path, *ext_arr):
    for file_ in os.listdir(dir_path):
        for ext in ext_arr:
            if fnmatch.fnmatch(file_, ext):
                yield file_
                break

def fuga(dir_path, *ext_arr):
    # 状況に応じて tuple や set, frozenset にするのもあり
    # 一番無難なのは list かな？
    return list(ifuga(dir_path, *ext_arr))

自分の site-packages より find 関数

自分はこんなのを site-packages の自作モジュールに仕込んでるので、こちらも公開してみる。

これならワイルドカードの代わりに正規表現オブジェクトや関数を渡しても動く。かわりに複数の条件はわたせない。複雑な条件は関数で何とかして、という方針。

正規表現を渡した時の動作は、個人的な好みで match を使っている。 search にしたい場合は、正規表現オブジェクトそのものではなく re.compile(u'python').search といったメソッドオブジェクトを渡してしまえばよい。

recursive 引数が真だと path の下位フォルダに潜っていく。なので、返す値もファイル名じゃなくてカレントからの相対パスもしくは絶対パスとして機能する文字列になっている。

今見直すと気になる点がちらほら。まず pattern の型チェックがきつすぎる気がした。正規表現オブジェクトかどうかではなくて、 match メソッドをもつオブジェクト、程度でよいとおもう。あと main 関数内の unicode(arg, sys.getfilesystemencoding()) 。 getfilesystemencoding じゃなくて getdefalutencoding だ。つまり unicode(arg) で十分だ。

# coding: utf-8

import fnmatch
import functools
import os
import re

def ifind(path, pattern, recursive=True):
    u"""path ディレクトリから条件に合うファイルを探して一つずつ返す


    pattern は u'*.py' のような Unix シェルライクのワイルドカードか、
    正規表現オブジェクトか、ファイル名を受け取って真偽値を返す関数。
    """
    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

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

def find(path, pattern, recursive=True):
    u"""path ディレクトリから条件に合うファイルを探す


    pattern は u'*.py' のような Unix シェルライクのワイルドカードか、
    正規表現オブジェクトか、ファイル名を受け取って真偽値を返す関数。
    """
    return list(ifind(path, pattern, recursive))

def main():
    import sys
    import optparse

    usage = u'usage: %prog dir pattern'
    parser = optparse.OptionParser(usage=usage)
    parser.add_option(
            '-m', '--move',
            action='store',
            dest='move',
            metavar='DIR',
            help=u'発見したファイルを DIR に移動します',
            )
    parser.set_defaults(move=None)
    options, args = parser.parse_args()

    if options.move:
        if not os.path.isdir(options.move):
            parser.error(u'%s is not directory' % options.move)


    args_len = len(args)
    if args_len != 2:
        parser.error(u'takes exactly 2 arguments (%d given)' % args_len)

    dir, pattern = [unicode(arg, sys.getfilesystemencoding()) for arg in args]

    if not os.path.isdir(dir):
        parser.error(u'%s is not directory' % dir)
    for path in ifind(dir, pattern):
        print path
        if options.move:
            newpath = os.path.join(options.move, os.path.basename(path))
            if path != newpath:
                if not os.path.isfile(newpath):
                    print u"moved: from %s\n       to %s" % (path, newpath)
                    os.rename(path, newpath)
                else:
                    print u"can not moved: %s" % path

if __name__ == '__main__':
    main()

複数の条件は渡せないので、複数の拡張子を探すにはもう一工夫。こう、かな。 exts には (u'.jpg', u'.png') のようなタプルを渡す方向で作成。

def ifindext(path, exts, recursive=True):
    u"""ある拡張子を持つファイルを探して一つずつ返す

    exts は (u'.py', u'.c', u'.h') のようなタプル。
    """
    return ifind(
            path, lambda ext: os.path.splitext(ext)[1] in exts, recursive)

def findext(path, exts, recursive=True):
    u"""ある拡張子を持つファイルを探す

    exts は (u'.py', u'.c', u'.h') のようなタプル。
    """
    return list(ifindext(path, exts, recursive))