MetaGPT

1
#!/usr/bin/env python
2
# -*- coding: utf-8 -*-
3
"""
4
@Time    : 2024/3/11
5
@Author  : mashenquan
6
@File    : tree.py
7
@Desc    : Implement the same functionality as the `tree` command.
8
        Example:
9
            >>> print_tree(".")
10
            utils
11
            +-- serialize.py
12
            +-- project_repo.py
13
            +-- tree.py
14
            +-- mmdc_playwright.py
15
            +-- cost_manager.py
16
            +-- __pycache__
17
            |   +-- __init__.cpython-39.pyc
18
            |   +-- redis.cpython-39.pyc
19
            |   +-- singleton.cpython-39.pyc
20
            |   +-- embedding.cpython-39.pyc
21
            |   +-- make_sk_kernel.cpython-39.pyc
22
            |   +-- file_repository.cpython-39.pyc
23
            +-- file.py
24
            +-- save_code.py
25
            +-- common.py
26
            +-- redis.py
27
"""
28
from __future__ import annotations
29

30
import subprocess
31
from pathlib import Path
32
from typing import Callable, Dict, List
33

34
from gitignore_parser import parse_gitignore
35

36

37
def tree(root: str | Path, gitignore: str | Path = None, run_command: bool = False) -> str:
38
    """
39
    Recursively traverses the directory structure and prints it out in a tree-like format.
40

41
    Args:
42
        root (str or Path): The root directory from which to start traversing.
43
        gitignore (str or Path): The filename of gitignore file.
44
        run_command (bool): Whether to execute `tree` command. Execute the `tree` command and return the result if True,
45
            otherwise execute python code instead.
46

47
    Returns:
48
        str: A string representation of the directory tree.
49

50
    Example:
51
            >>> tree(".")
52
            utils
53
            +-- serialize.py
54
            +-- project_repo.py
55
            +-- tree.py
56
            +-- mmdc_playwright.py
57
            +-- __pycache__
58
            |   +-- __init__.cpython-39.pyc
59
            |   +-- redis.cpython-39.pyc
60
            |   +-- singleton.cpython-39.pyc
61
            +-- parse_docstring.py
62

63
            >>> tree(".", gitignore="../../.gitignore")
64
            utils
65
            +-- serialize.py
66
            +-- project_repo.py
67
            +-- tree.py
68
            +-- mmdc_playwright.py
69
            +-- parse_docstring.py
70

71
            >>> tree(".", gitignore="../../.gitignore", run_command=True)
72
            utils
73
            ├── serialize.py
74
            ├── project_repo.py
75
            ├── tree.py
76
            ├── mmdc_playwright.py
77
            └── parse_docstring.py
78

79

80
    """
81
    root = Path(root).resolve()
82
    if run_command:
83
        return _execute_tree(root, gitignore)
84

85
    git_ignore_rules = parse_gitignore(gitignore) if gitignore else None
86
    dir_ = {root.name: _list_children(root=root, git_ignore_rules=git_ignore_rules)}
87
    v = _print_tree(dir_)
88
    return "\n".join(v)
89

90

91
def _list_children(root: Path, git_ignore_rules: Callable) -> Dict[str, Dict]:
92
    dir_ = {}
93
    for i in root.iterdir():
94
        if git_ignore_rules and git_ignore_rules(str(i)):
95
            continue
96
        try:
97
            if i.is_file():
98
                dir_[i.name] = {}
99
            else:
100
                dir_[i.name] = _list_children(root=i, git_ignore_rules=git_ignore_rules)
101
        except (FileNotFoundError, PermissionError, OSError):
102
            dir_[i.name] = {}
103
    return dir_
104

105

106
def _print_tree(dir_: Dict[str:Dict]) -> List[str]:
107
    ret = []
108
    for name, children in dir_.items():
109
        ret.append(name)
110
        if not children:
111
            continue
112
        lines = _print_tree(children)
113
        for j, v in enumerate(lines):
114
            if v[0] not in ["+", " ", "|"]:
115
                ret = _add_line(ret)
116
                row = f"+-- {v}"
117
            else:
118
                row = f"    {v}"
119
            ret.append(row)
120
    return ret
121

122

123
def _add_line(rows: List[str]) -> List[str]:
124
    for i in range(len(rows) - 1, -1, -1):
125
        v = rows[i]
126
        if v[0] != " ":
127
            return rows
128
        rows[i] = "|" + v[1:]
129
    return rows
130

131

132
def _execute_tree(root: Path, gitignore: str | Path) -> str:
133
    args = ["--gitfile", str(gitignore)] if gitignore else []
134
    try:
135
        result = subprocess.run(["tree"] + args + [str(root)], capture_output=True, text=True, check=True)
136
        if result.returncode != 0:
137
            raise ValueError(f"tree exits with code {result.returncode}")
138
        return result.stdout
139
    except subprocess.CalledProcessError as e:
140
        raise e
141