自定义TreeNode类递归构建Python目录树的层级结构问题

问题分析 #

原代码的核心问题在于build_tree函数没有维护节点的层级上下文：

• 依赖全局变量tree，所有新创建的节点都直接添加到根节点，无法关联到对应的父目录节点。
• os.walk的遍历逻辑与TreeNode的层级构建完全脱节，没有根据目录路径找到对应的父节点。
• 递归分支的逻辑混乱，working_directories变量无法正确传递层级信息。

修正方案 #

重构build_tree函数，让它接收当前父节点和当前目录路径作为参数，递归时将子节点挂载到正确的父节点上。同时统一类名大小写（原代码Treenode改为TreeNode，符合Python命名规范）。

修正后的完整代码 #

import os

class TreeNode:
    def __init__(self, data):
        self.data = data
        self.children = []
        self.parent = None

    def add_child(self, child):
        child.parent = self
        self.children.append(child)

    def get_level(self):
        level = 0
        p = self.parent
        while p:
            level += 1
            p = p.parent
        return level

    def print_tree(self):
        spaces = " " * self.get_level() * 3
        prefix = spaces + "|__" if self.parent else ""
        print(prefix + self.data)
        for child in self.children:
            child.print_tree()

def build_tree(parent_node, dir_path):
    # 获取当前目录下的所有条目（文件+子目录），按目录在前、文件在后排序（匹配期望输出结构）
    entries = sorted(os.listdir(dir_path), key=lambda x: (not os.path.isdir(os.path.join(dir_path, x)), x))
    for entry in entries:
        entry_path = os.path.join(dir_path, entry)
        # 创建当前条目的节点
        node = TreeNode(entry)
        parent_node.add_child(node)
        # 如果是目录，递归构建其子节点
        if os.path.isdir(entry_path):
            build_tree(node, entry_path)

if __name__ == '__main__':
    root_dir = "C:/Level0"
    # 创建根节点，显示完整路径
    tree = TreeNode(root_dir)
    build_tree(tree, root_dir)
    tree.print_tree()

逻辑说明 #

1. 层级上下文传递：build_tree每次递归时，都会将当前目录对应的节点作为父节点传入，确保子目录/文件节点被挂载到正确的父节点下。
2. 条目排序：使用sorted函数将目录排在文件前面，和期望输出的结构完全匹配（若不需要排序可去掉该逻辑）。
3. 递归触发：仅当遍历到目录时才递归调用build_tree，文件节点不会触发递归，完全符合目录树的结构逻辑。

输出效果 #

修正后的代码会输出符合期望的层级结构：

C:/Level0
   |__Level1
      |__Level2
         |__Level3
            |__LEvel4
               |__4file1.txt
            |__3file1.txt
         |__2file1.txt
      |__Level2b
         |__2bfile1.txt
      |__1file1.txt
      |__1file2.txt
   |__0file.txt

内容的提问来源于stack exchange，提问作者user20428297