python实现输出文件夹的目录树结构

最近清理QQ文件夹的时候, 发现群聊图片巨TM大, 就想看看群友们都发了什么怪图片
但是QQ群聊图片文件夹里面的文件夹是一个套一个, 靠手工copy肯定是不太行的, 就写了个简单的脚本遍历里面所有的文件

之后在用tree命令的时候想了想, 自己能不能也简单实现下类似tree的目录的树状结构的打印

我们来一步一步实现一下

遍历文件夹内的所有文件#

首先简单递归遍历一个文件夹

import os

def get_all_file(path):
    files = []
    items = os.listdir(path)
    for i in items:
        i_path = path + "/" + i
        if os.path.isdir(i_path):
            files.extend(get_all_file(i_path))
        else:
            files.append(i_path)
    return files

分析tree命令的输出#

输出树状结构可以把文件目录按树数据结构进行保存打印输出

不过既然遍历了所有文件和文件夹, 我本身通过就可以判断当前的层数来输出对应的树干和树枝

我新建了一个文件夹, 在里面嵌套了一些文件夹和文件, 先简单看下tree命令是怎么输出这个文件夹的结构的

├── 1
│  ├── 11
│  ├── 12
│  └── 13
├── 1.txt
├── 2
│  ├── 21
│  └── 22
│     ├── 3
│     │  └── 1.txt
│     └── 3 - 副本
│        └── 1.txt
└── 3
   ├── 1.txt
   └── 31
      └── 1.txt

输出树的形状用了 │ ├ └ ─ 四种符号来表示树干, 树枝分支 和 最后一个树枝

分析一下分支形状输出的逻辑

1. 当前层不是最后一枝, 那么就使用 ├ 作为分支形状, 是最后一枝则用 └ 作为分支形状

2. ├ 和 ├ └之间要用 │ 进行连接, └ 之后则不需要 │

其中第2点可以这样判断, 判断父节点是不是最后一枝, 不是则输出 │ 来连接树干

具体实现#

通过当前层级所在的目录列表的下标和列表长度判断当前文件/文件夹是不是最后一枝

然后需要一个参数depth表示当前层数, 递归的时候depth=depth+1来表示层数增加

def print_folder_tree(path, depth=0):
    files = []
    items = os.listdir(path)
    for index, i in enumerate(items):
        # 是否是最后一个元素
        is_last = index == len(items) - 1
        # 拼接文件路径
        i_path = path + "/" + i
        # 根据层数打印空格
        print("   " * depth,end="")
        if is_last:
            print("└── ", end="")
        else:
            print("├── ", end="")
        # 如果是文件夹, 递归
        if os.path.isdir(i_path):
            print(i)
            files.extend(print_folder_tree(path=i_path, depth=depth + 1))
        # 如果是文件就把路径添加到files数组
        else:
            print(i_path.split("/")[-1])
            files.append(i_path)
    return files  

代码输出的结果是这样, 可以看到层级显示正常

├── 1
   ├── 11
   ├── 12
   └── 13
├── 1.txt
├── 2
   ├── 21
   └── 22
      ├── 3
         └── 1.txt
      └── 3 - 副本
         └── 1.txt
└── 3
   ├── 1.txt
   └── 31
      └── 1.txt

接下来需要输出树干 │ , 但是树干的输出需要根据父节点, 父节点的父节点...是否是所在层级的最后一枝来输出

├── 2
│	├── 21
│	└── 22
│		├── 3
│		│	└── 1.txt
│		└── 3 - 副本
│			└── 1.txt

我这里使用parent_is_last, 来1和0表示每个父节点是否为最后一支, 类似于独热码(One Hot Encode)
同时parent_is_last的长度也可以表示目录的层数

以 22/3/1.txt 为例, 所在层级是3, parent_is_last就是1010, 首位1忽略, 实际根据010输出树干, 所以它往上三层父节点的 │ 树干就是 有 没有 有

然后加上了控制遍历深度和空格宽度的参数

具体实现如下

def print_folder_tree(path, parent_is_last=1, depth_limit=-1, tab_width=1):
    """
    以树状打印输出文件夹下的文件, 并返回文件夹内的所有文件
    :param tab_width: 空格宽度
    :param path: 文件夹路径
    :param depth_limit: 要输出文件夹的层数, -1为输出全部文件及文件夹
    :param parent_is_last: 递归调用上级文件夹是否是最后一个文件(夹), 控制输出 │ 树干
    :return: 返回path下的所有文件的数组
    """
    files = []
    if len(str(parent_is_last)) - 1 == depth_limit:
        return files
    items = os.listdir(path)
    for index, i in enumerate(items):
        is_last = index == len(items) - 1
        i_path = path + "/" + i
        for k in str(parent_is_last)[1:]:
            if k == "0":
                print("│" + "\t" * tab_width, end="")
            if k == "1":
                print("\t" * tab_width, end="")
        if is_last:
            print("└── ", end="")
        else:
            print("├── ", end="")
        if os.path.isdir(i_path):
            print(i)
            files.extend(print_folder_tree(
                path=i_path, depth_limit=depth_limit, parent_is_last=(parent_is_last * 10 + 1) if is_last else (parent_is_last * 10)))
        else:
            print(i_path.split("/")[-1])
            files.append(i_path)
    return files

最后代码的运行结果是这样的

├── 1
│	├── 11
│	├── 12
│	└── 13
├── 1.txt
├── 2
│	├── 21
│	└── 22
│		├── 3
│		│	└── 1.txt
│		└── 3 - 副本
│			└── 1.txt
└── 3
	├── 1.txt
	└── 31
		└── 1.txt