torch.package.file_structure_representation 的源代码

from typing import Dict, List

from .glob_group import GlobGroup, GlobPattern

__all__ = ["Directory"]


class Directory:
    """文件结构的表示。组织为具有其目录子列表的目录节点。通过调用
    :meth:`PackageImporter.file_structure`来创建包的目录。"""

    def __init__(self, name: str, is_dir: bool):
        self.name = name
        self.is_dir = is_dir
        self.children: Dict[str, Directory] = {}

    def _get_dir(self, dirs: List[str]) -> "Directory":
        """如果尚未构建，则构建目录路径并返回列表中的最后一个目录
        。

        参数:
            dirs (List[str]): 被视为路径的目录名称列表。

        返回:
            :class:`Directory`: dirs列表中指定的最后一个目录。
        """
        if len(dirs) == 0:
            return self
        dir_name = dirs[0]
        if dir_name not in self.children:
            self.children[dir_name] = Directory(dir_name, True)
        return self.children[dir_name]._get_dir(dirs[1:])

    def _add_file(self, file_path: str):
        """向目录添加文件。

        参数:
            file_path (str): 要添加的文件路径。最后一个元素作为文件添加，
                其他路径项作为目录添加。
        """
        *dirs, file = file_path.split("/")
        dir = self._get_dir(dirs)
        dir.children[file] = Directory(file, False)

    def has_file(self, filename: str) -> bool:
        """检查文件是否存在于:class:`Directory`中。

        参数:
            filename (str): 要搜索的文件路径。
        返回:
            bool: 如果:class:`Directory`包含指定的文件。
        """
        lineage = filename.split("/", maxsplit=1)
        child = lineage[0]
        grandchildren = lineage[1] if len(lineage) > 1 else None
        if child in self.children.keys():
            if grandchildren is None:
                return True
            else:
                return self.children[child].has_file(grandchildren)
        return False

    def __str__(self):
        str_list: List[str] = []
        self._stringify_tree(str_list)
        return "".join(str_list)

    def _stringify_tree(
        self, str_list: List[str], preamble: str = "", dir_ptr: str = "─── "
    ):
        """递归方法生成目录的打印友好版本。"""
        space = "    "
        branch = "│   "
        tee = "├── "
        last = "└── "

        # 添加此目录的表示
        str_list.append(f"{preamble}{dir_ptr}{self.name}\n")

        # 添加目录的子表示
        if dir_ptr == tee:
            preamble = preamble + branch
        else:
            preamble = preamble + space

        file_keys: List[str] = []
        dir_keys: List[str] = []
        for key, val in self.children.items():
            if val.is_dir:
                dir_keys.append(key)
            else:
                file_keys.append</