pytorch

1
from typing import Dict, List
2

3
from .glob_group import GlobGroup, GlobPattern
4

5
__all__ = ["Directory"]
6

7

8
class Directory:
9
    """A file structure representation. Organized as Directory nodes that have lists of
10
    their Directory children. Directories for a package are created by calling
11
    :meth:`PackageImporter.file_structure`."""
12

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

18
    def _get_dir(self, dirs: List[str]) -> "Directory":
19
        """Builds path of Directories if not yet built and returns last directory
20
        in list.
21

22
        Args:
23
            dirs (List[str]): List of directory names that are treated like a path.
24

25
        Returns:
26
            :class:`Directory`: The last Directory specified in the dirs list.
27
        """
28
        if len(dirs) == 0:
29
            return self
30
        dir_name = dirs[0]
31
        if dir_name not in self.children:
32
            self.children[dir_name] = Directory(dir_name, True)
33
        return self.children[dir_name]._get_dir(dirs[1:])
34

35
    def _add_file(self, file_path: str):
36
        """Adds a file to a Directory.
37

38
        Args:
39
            file_path (str): Path of file to add. Last element is added as a file while
40
                other paths items are added as directories.
41
        """
42
        *dirs, file = file_path.split("/")
43
        dir = self._get_dir(dirs)
44
        dir.children[file] = Directory(file, False)
45

46
    def has_file(self, filename: str) -> bool:
47
        """Checks if a file is present in a :class:`Directory`.
48

49
        Args:
50
            filename (str): Path of file to search for.
51
        Returns:
52
            bool: If a :class:`Directory` contains the specified file.
53
        """
54
        lineage = filename.split("/", maxsplit=1)
55
        child = lineage[0]
56
        grandchildren = lineage[1] if len(lineage) > 1 else None
57
        if child in self.children.keys():
58
            if grandchildren is None:
59
                return True
60
            else:
61
                return self.children[child].has_file(grandchildren)
62
        return False
63

64
    def __str__(self):
65
        str_list: List[str] = []
66
        self._stringify_tree(str_list)
67
        return "".join(str_list)
68

69
    def _stringify_tree(
70
        self, str_list: List[str], preamble: str = "", dir_ptr: str = "─── "
71
    ):
72
        """Recursive method to generate print-friendly version of a Directory."""
73
        space = "    "
74
        branch = "│   "
75
        tee = "├── "
76
        last = "└── "
77

78
        # add this directory's representation
79
        str_list.append(f"{preamble}{dir_ptr}{self.name}\n")
80

81
        # add directory's children representations
82
        if dir_ptr == tee:
83
            preamble = preamble + branch
84
        else:
85
            preamble = preamble + space
86

87
        file_keys: List[str] = []
88
        dir_keys: List[str] = []
89
        for key, val in self.children.items():
90
            if val.is_dir:
91
                dir_keys.append(key)
92
            else:
93
                file_keys.append(key)
94

95
        for index, key in enumerate(sorted(dir_keys)):
96
            if (index == len(dir_keys) - 1) and len(file_keys) == 0:
97
                self.children[key]._stringify_tree(str_list, preamble, last)
98
            else:
99
                self.children[key]._stringify_tree(str_list, preamble, tee)
100
        for index, file in enumerate(sorted(file_keys)):
101
            pointer = last if (index == len(file_keys) - 1) else tee
102
            str_list.append(f"{preamble}{pointer}{file}\n")
103

104

105
def _create_directory_from_file_list(
106
    filename: str,
107
    file_list: List[str],
108
    include: "GlobPattern" = "**",
109
    exclude: "GlobPattern" = (),
110
) -> Directory:
111
    """Return a :class:`Directory` file structure representation created from a list of files.
112

113
    Args:
114
        filename (str): The name given to the top-level directory that will be the
115
            relative root for all file paths found in the file_list.
116

117
        file_list (List[str]): List of files to add to the top-level directory.
118

119
        include (Union[List[str], str]): An optional pattern that limits what is included from the file_list to
120
            files whose name matches the pattern.
121

122
        exclude (Union[List[str], str]): An optional pattern that excludes files whose name match the pattern.
123

124
    Returns:
125
            :class:`Directory`: a :class:`Directory` file structure representation created from a list of files.
126
    """
127
    glob_pattern = GlobGroup(include, exclude=exclude, separator="/")
128

129
    top_dir = Directory(filename, True)
130
    for file in file_list:
131
        if glob_pattern.matches(file):
132
            top_dir._add_file(file)
133
    return top_dir
134