[Python]フォルダ構造をDataFrameに入れる方法

はじめに

フォルダ毎や形式毎にファイルを読み込んで処理する時、それぞれのファイルpathを検索しながら取得するのは面倒ではありませんか？
そこで、ファイル情報をDataFrameに格納することで、直感的にpathを取得出来るようにしました。

サンプル

以下のようなディレクトリ構成があったとき、

tree構造

/data
└── dir_A
    └── a.txt
└── dir_B
    └── b.txt

以下のようなDataFrameを取得します。

index	ファイル名	相対path	ファイル形式	a	b	c	絶対path
0	data	./data	ディレクトリ	data			C:/data
1	dir_A	./data/dir_A	ディレクトリ	data	dir_A		C:/data/dir_A
2	a.txt	./data/dir_A/a.txt	.txt	data	dir_A	a.txt	C:/data/dir_A
3	dir_B	./data/dir_B	ディレクトリ	data	dir_B		C:/data/dir_B
4	b.txt	./data/dir_B/b.txt	.txt	data	dir_B	b.txt	C:/data/dir_B

コード

tree_to_df

import pathlib
import pandas as pd
import string

def tree_to_df(path ='.', dir_only=False, file_only=False, path_only=False):

    # get target_path
    p = pathlib.Path(path)

    # get tree
    pathlist = list(p.glob("**/*"))

    # get file_info
    namelist = [i.name for i in pathlist]
    str_p = p.resolve().as_posix()
    str_pathlist = [i.resolve().as_posix() for i in pathlist]
    str_relpath = [path + '/' + i.replace(str_p + '/', '') for i in str_pathlist]
    str_pathlist_split = [i.replace(str_p + '/', '').split('/') for i in str_pathlist]
    type_list = [i.suffix if i.suffix !='' else 'dir'  for i in pathlist]

    def col_name(num):
        n = map(int, str(num))
        w = [string.ascii_lowercase[i] for i in n]
        result = ''.join(w)
        return result

    # make_df
    df_splitpath = pd.DataFrame(str_pathlist_split)
    df_splitpath = df_splitpath.rename(columns=lambda s: col_name(s))
    df_namelist = pd.DataFrame(namelist, columns= ['name'])
    df_type_list = pd.DataFrame(type_list, columns= ['type'])
    df_str_relpath = pd.DataFrame(str_relpath, columns= ['relpath'])
    df_str_abspath = pd.DataFrame(str_pathlist, columns= ['abspath'])

    # result_df
    df = pd.concat([df_namelist,df_str_relpath,df_type_list,df_splitpath,df_str_abspath],axis=1)

    # select_df
    df = df = df[df['type']=='dir'] if dir_only else df
    df = df[df['type']!='dir'] if file_only else df
    df = df['abspath'] if path_only else df

    return df

使い方

検索したいpathを関数に渡します。引数で、フォルダ・ファイル・pathのみ出力もできます。出力したDataFrameにqueryを叩くことで、直感的にpathを取得することが出来ます。

use

# cwd
tree_to_df()
# cwd
tree_to_df('.')
# up_dir
tree_to_df('../')

参考文献

Pathlibチートシート - Qiita
Pythonでフォルダ内のファイルリストを取得する - Qiita
Python, pathlibでファイル名・拡張子・親ディレクトリを取得 | note.nkmk.me
Pythonの三項演算子（条件演算子）でif文を一行で書く | note.nkmk.me
string - Alphabet range python - Stack Overflow
[python]リストの中身を文字列に変換する。 | 岡田デザインラボ