treeコマンドもどき用のジェネレータを作ってみた

はじめに

• Pythonでコマンドプロンプトのtreeっぽいことをやりたくなった。
• 見た目の調整部分は文字幅調整などが面倒そうなので放っておいてとりあえずフォルダやファイルの名前と深さ、ファイルかフォルダのどちらかの情報を返すジェネレータとして作ってみた。
• GUIで出力するときも階層と名前だけの出力があれば使えるので応用すれば他にも使えそうなものができた。

できたもののコード

from pathlib import Path
from collections.abc import Generator


def tree(path: Path, nameonly: bool = True, root_nameonly: bool = True) -> Generator[tuple[int, str, bool], None, None]:
    if nameonly or root_nameonly:
        yield len(path.parents), path.name, path.is_dir()
    else:
        yield len(path.parents), str(path), path.is_dir()
    for p in sorted(path.iterdir(), key=Path.is_file):
        if p.is_dir():
            yield from tree(p, nameonly=nameonly, root_nameonly=True)
        else:
            yield len(p.parents), p.name if nameonly else str(p), p.is_dir()


for depth, path, is_dir in tree(Path()):
    print(f"{'  ' * depth}{path}[{'D' if is_dir else 'F'}]")

引数の説明

今回作成したtree関数はフォルダをpathlib.Pathとして受け取り、そこを起点に探索を行うジェネレータ関数。
引数のnameonlyはyieldで戻ってくるものを最初に与えたpathからの名前だけにするかを決める。
root_nameonlyは最初に与えたpathだけを名前だけにするかを決める。
フルパスを与えて、一番最初に表示するときはフルパス、以降は名前だけということもできる。

戻り値の説明

戻り値は深さ(int)、パスまたは名前(str)、フォルダかどうか(bool)の3要素からなるタプル。
深さはlen(path.parents)を使って取得しているので最初に与えたフォルダが1となる。
深さを * でインデント用のスペースでかけるだけでもインデントで分けられたツリーっぽいものが簡単に作れる。
フォルダかどうかの情報もあるので、ファイルとフォルダで見た目になにか工夫を加えたい場合はそこを使うことができる。

※アノテーションを付けるとき、collections.abcのGeneratorをアノテーションに使えないバージョンではtypingの方を使用する。

実装について

全体的にpathlib.Pathの機能を使っているだけなので思った以上にシンプルに作れている。
地味にやってよかったと思っているのはsortedを使ったところ。
path.iterdir()をするときにsortedを挟むことで最初にフォルダが来るようにしている。
上のコードではkeyに渡す関数にpathlib.Pathクラスのis_fileを渡している。
クラスの普通のメソッドでもクラスから直接取ると第一引数にそのクラスのインスタンスを要求する形になるので、path.iterdir()でやってくるpathlib.Pathのインスタンスに対して正しく動いてくれる。
無名関数を使って.is_file()をする下のような形でも同じ動きをしてくれる。

# どっちを使っても同じ動きになる
sorted(path.iterdir(), key=Path.is_file)
sorted(path.iterdir(), key=lambda p: p.is_file())

判定にis_fileを使っているので、フォルダだった場合False(0)が返るのでデフォルトの昇順でソートされればフォルダ一番最初に来るようになる。

最初に引数として渡されたフォルダの親までの数と、パスか名前、フォルダかどうかの真偽値をまとめてyieldで返す。
あとは巡回した要素がフォルダかファイルかを判定して、フォルダならyield fromで再帰する。
ファイルの場合は親までの数と、パスか名前、フォルダかどうかの真偽値をまとめてyieldで返す。

終わりに

pathlibがとにかく使いやすい。
strを噛ませないと引数として渡せないことがたまにあるが、それを差し引いてもパス関連は全てこれで済むというところがわかりやすい。
内部の処理も元となるフォルダをどうするかと、その中身を巡回して各要素に対してそれぞれどうするかというシンプルな構造で作ることができた。