用Python分分钟造个Tree命令玩一下

前言

在处理大量文件和目录时，清晰地了解它们的结构是非常重要的。在命令行中，我们经常使用tree命令来以树状图的形式显示文件和目录结构。然而，有时我们可能需要在 Python 中进行相同的操作，以便更好地处理数据、生成报告或进行自动化任务。

Python 的tree命令：

Python 的强大性质和丰富的标准库使得实现tree命令变得相对简单。我们可以使用 Python 的pathlib模块来轻松访问文件系统，以及click模块来处理命令行参数。这使得编写自定义的tree命令变得非常容易。

功能亮点：

• 目录可视化： Python实现的tree命令允许您以直观的树状图形式可视化文件和目录结构，使其易于理解和分析。

• 深度限制： 您可以设置深度限制，只显示指定层级以下的文件和目录。这对于快速查看结构非常有用。

• 隐藏文件控制： 您可以选择是否显示隐藏文件，以便在需要时包含或排除它们。

• 参数缩写： 使用命令行参数的缩写形式，更快速地配置tree命令的行为。

• 自定义颜色： 使用 ANSI 颜色代码，为文件和目录添加自定义颜色，以便更好地突出显示它们。

用途：

Python 实现的tree命令在许多领域都有广泛的用途，包括：

• 开发： 在软件开发中，可视化项目目录结构有助于了解代码库的组织，查找特定文件和快速导航。

• 数据分析： 对于数据分析师和科学家，可视化数据文件和目录结构有助于管理和处理数据集。

• 自动化任务： 通过编写脚本，您可以自动执行文件操作、备份文件、或创建报告，而无需手动查看目录结构。

• 教育： Python "tree" 命令还可用于教育目的，帮助学生更好地理解文件系统和目录结构的概念。

用到的符号实际只有三个├─、│、└─.

列出当前路径的上一层路径下的目录树

#! -*-conding=: UTF-8 -*-
# 2023/9/25 16:50
"""
    描述：python实现tree命令
"""

from pathlib import Path
from typing import List

# 黑名单 不想列出的目录
blacklist: List[str] = [".git", ".idea"]
# 指定列出哪个目录下的所有内容，使用当前路径的上一层路径
rpath: Path = Path(__file__).resolve().parent


def dir_and_file(path: Path, symbol: str = ""):
    # 列出所有目录和文件 同时统计数量用于判断
    file_list = [entry for entry in path.iterdir() if entry.name not in blacklist]
    total_num = len(file_list)
    num = 1
    for entry in file_list:
        # 路径合并 递归调用时继续向下传递
        if entry.is_file():
            # 判断是否为最后一个
            prefix = "  └─ " if num == total_num else "  ├─ "
            print(f"{symbol}{prefix}{entry.name}")
        else:
            # 判断目录是否为最后一个 如果是则使用不同的符号
            prefix = "  └─ " if num == total_num else "  ├─ "
            print(f"{symbol}{prefix}{entry.name}")
            dir_and_file(entry, f"{symbol}{'     ' if num == total_num else '  │  '}")
        num += 1


if __name__ == '__main__':
    print(rpath)
    dir_and_file(rpath)

限制显示的目录树最大深度

#! -*-conding=: UTF-8 -*-
# 2023/9/25 16:50
"""
    描述：python实现tree命令
"""

from pathlib import Path
from typing import List

# 黑名单 不想列出的目录
blacklist: List[str] = [".git", ".idea"]
# 指定列出哪个目录下的所有内容，使用当前路径的上一层路径
rpath: Path = Path(__file__).resolve().parent


def dir_and_file(path: Path, symbol: str = "", max_depth: int = float('inf'), current_depth: int = 1):
    if current_depth > max_depth:
        return

    # 列出所有目录和文件 同时统计数量用于判断
    file_list = [entry for entry in path.iterdir() if entry.name not in blacklist]
    total_num = len(file_list)
    num = 1
    for entry in file_list:
        # 路径合并 递归调用时继续向下传递
        if entry.is_file():
            # 判断是否为最后一个
            prefix = "  └─ " if num == total_num else "  ├─ "
            print(f"{symbol}{prefix}{entry.name}")
        else:
            # 判断目录是否为最后一个 如果是则使用不同的符号
            prefix = "  └─ " if num == total_num else "  ├─ "
            print(f"{symbol}{prefix}{entry.name}")
            dir_and_file(entry, f"{symbol}{'     ' if num == total_num else '  │  '}", max_depth, current_depth + 1)
        num += 1


if __name__ == '__main__':
    print(rpath)
    depth = 1  # 设置显示的深度
    dir_and_file(rpath, max_depth=depth)

给目录上点颜色

#! -*-conding=: UTF-8 -*-
# 2023/9/25 16:50
"""
    描述：python实现tree命令
"""

from pathlib import Path
from typing import List

# 黑名单 不想列出的目录
blacklist: List[str] = [".git", ".idea"]
# 指定列出哪个目录下的所有内容，使用当前路径的上一层路径
rpath: Path = Path(__file__).resolve().parent

# ANSI颜色代码
BLUE = "\033[34m"
WHITE = "\033[0m"


def dir_and_file(path: Path, symbol: str = "", max_depth: int = float('inf'), current_depth: int = 1):
    if current_depth > max_depth:
        return

    # 列出所有目录和文件 同时统计数量用于判断
    file_list = [entry for entry in path.iterdir() if entry.name not in blacklist]
    total_num = len(file_list)
    num = 1
    for entry in file_list:
        # 路径合并 递归调用时继续向下传递
        if entry.is_file():
            # 判断是否为最后一个
            prefix = "  └─ " if num == total_num else "  ├─ "
            print(f"{symbol}{prefix}{WHITE}{entry.name}")
        else:
            # 判断目录是否为最后一个 如果是则使用不同的符号
            prefix = "  └─ " if num == total_num else "  ├─ "
            print(f"{symbol}{prefix}{BLUE}{entry.name}{WHITE}")
            dir_and_file(entry, f"{symbol}{'     ' if num == total_num else '  │  '}", max_depth, current_depth + 1)
        num += 1


if __name__ == '__main__':
    print(rpath)
    depth = 1  # 设置显示的深度
    dir_and_file(rpath, max_depth=depth)

引入click模块实现命令行传参

首先需要安装click模块

pip install click

使用click实现命令行传参

#! -*-conding=: UTF-8 -*-
# 2023/9/25 16:50
"""
    描述：python实现tree命令
"""

from pathlib import Path
import click
from typing import List

# ANSI颜色代码
BLUE = "\033[34m"
WHITE = "\033[0m"


def is_hidden(file: Path) -> bool:
    return file.name.startswith(".")


def get_entries(path: Path, show_hidden: bool = False) -> List[Path]:
    entries = list(path.iterdir())
    if not show_hidden:
        entries = [entry for entry in entries if not is_hidden(entry)]
    return sorted(entries, key=lambda entry: entry.is_file())


def print_tree(path: Path, symbol: str = "", max_depth: int = float("inf"), current_depth: int = 1,
               show_hidden: bool = False):
    if current_depth > max_depth:
        return

    entries = get_entries(path, show_hidden)
    total_entries = len(entries)
    num = 1
    for entry in entries:
        is_dir = entry.is_dir()
        entry_name = entry.name

        if num == total_entries:
            branch = "└─ "
            new_symbol = symbol + "     " if num == total_entries else symbol + "  │  "
        else:
            branch = "├─ "
            new_symbol = symbol + "  │  "

        if is_dir:
            click.echo(f"{symbol}{branch}{BLUE}{entry_name}{WHITE}")
            print_tree(entry, new_symbol, max_depth, current_depth + 1, show_hidden)
        else:
            click.echo(f"{symbol}{branch}{WHITE}{entry_name}")

        num += 1


@click.command()
@click.option("-p", "--path", default=".", type=click.Path(exists=True), help="要列出内容的目录路径")
@click.option("-d", "--depth", default=1, type=int, help="限制显示的深度")
@click.option("--hidden", is_flag=True, help="显示隐藏文件")
def main(path, depth, hidden):
    path = Path(path).resolve()
    click.echo(path)
    print_tree(path, max_depth=depth, show_hidden=hidden)


if __name__ == '__main__':
    main()

使用内置的argparse模块实现命令行传参

因为click模块需要单独安装，有的时候不是太方便，可以使用内置的argparse模块代替。

#! -*-conding=: UTF-8 -*-
# 2023/9/25 16:50
"""
    描述：python实现tree命令
"""

from pathlib import Path
import argparse

from typing import Optional

# 黑名单 不想列出的目录
blacklist = [".git", ".idea"]

# ANSI颜色代码
BLUE = "\033[34m"
WHITE = "\033[0m"


def is_hidden(file: Path) -> bool:
    return file.name.startswith(".")


def get_entries(path: Path, show_hidden: bool = False):
    entries = list(path.iterdir())
    if not show_hidden:
        entries = [entry for entry in entries if not is_hidden(entry)]
    return sorted(entries, key=lambda entry: entry.is_file())


def print_tree(path: Path, depth: Optional[int], show_hidden: bool = False, indent: str = "", is_last: bool = True):
    entries = get_entries(path, show_hidden)
    total_entries = len(entries)
    for i, entry in enumerate(entries):
        is_dir = entry.is_dir()
        entry_name = entry.name

        if i == total_entries - 1:
            branch = "└─ "
            new_indent = indent + "    " if is_last else indent + "│   "
        else:
            branch = "├─ "
            new_indent = indent + "│   "

        if is_dir:
            print(f"{indent}{branch}{BLUE}{entry_name}{WHITE}")
            if depth is None or (depth and depth > 1):
                print_tree(entry, None if depth is None else depth - 1, show_hidden, new_indent, i == total_entries - 1)
        else:
            print(f"{indent}{branch}{entry_name}")


def main():
    parser = argparse.ArgumentParser(description="Python实现tree命令")
    parser.add_argument("-p", "--path", type=str, nargs="?", default=".", help="要遍历的目录路径")
    parser.add_argument("-d", "--depth", type=int, help="限制显示的层数")
    parser.add_argument("--hidden", action="store_true", help="显示隐藏文件")
    args = parser.parse_args()
    path = Path(args.path).resolve()
    print(path)
    print_tree(path, args.depth, args.hidden)


if __name__ == "__main__":
    main()

shell中的tree命令参数还是有一些的，这里我们只实现部分基本参数相关的操作，感兴趣的小伙伴可以自行扩展。

喜欢这篇文章的话，就点个关注吧，或者关注一下我的公众号『海哥python』也可以，会持续分享高质量Python文章，以及其它内容。