From 684730e94f7c5fe7f5bac8a31ed8fac1937b6d43 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=8E=98=E6=9D=83=20=E9=A9=AC?= <mashenquan@fuzhi.ai>
Date: Wed, 13 Mar 2024 15:06:40 +0800
Subject: [PATCH] feat: +`tree` command

---
 metagpt/utils/tree.py            | 36 ++++++++++++++++++++++++++++----
 tests/metagpt/utils/test_tree.py | 16 +++++++++++---
 2 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/metagpt/utils/tree.py b/metagpt/utils/tree.py
index 1c0060842..c0386d822 100644
--- a/metagpt/utils/tree.py
+++ b/metagpt/utils/tree.py
@@ -27,17 +27,22 @@
 """
 from __future__ import annotations
 
+import subprocess
 from pathlib import Path
 from typing import Callable, Dict, List
 
+from gitignore_parser import parse_gitignore
 
-def tree(root: str | Path, git_ignore_rules: Callable = None) -> str:
+
+def tree(root: str | Path, gitignore: str | Path = None, run_command: bool = True) -> str:
     """
     Recursively traverses the directory structure and prints it out in a tree-like format.
 
     Args:
         root (str or Path): The root directory from which to start traversing.
-        git_ignore_rules (Callable): Optional. A function to filter files to ignore.
+        gitignore (str or Path): The filename of gitignore file.
+        run_command (bool): Whether to execute `tree` command. Execute the `tree` command and return the result if True,
+            otherwise execute python code instead.
 
     Returns:
         str: A string representation of the directory tree.
@@ -55,8 +60,7 @@ def tree(root: str | Path, git_ignore_rules: Callable = None) -> str:
             |   +-- singleton.cpython-39.pyc
             +-- parse_docstring.py
 
-            >>> from gitignore_parser import parse_gitignore
-            >>> tree(".", git_ignore_rules=parse_gitignore(full_path="../../.gitignore"))
+            >>> tree(".", gitignore="../../.gitignore")
             utils
             +-- serialize.py
             +-- project_repo.py
@@ -64,8 +68,21 @@ def tree(root: str | Path, git_ignore_rules: Callable = None) -> str:
             +-- mmdc_playwright.py
             +-- parse_docstring.py
 
+            >>> tree(".", gitignore="../../.gitignore", run_command=True)
+            utils
+            ├── serialize.py
+            ├── project_repo.py
+            ├── tree.py
+            ├── mmdc_playwright.py
+            └── parse_docstring.py
+
+
     """
     root = Path(root).resolve()
+    if run_command:
+        return _execute_tree(root, gitignore)
+
+    git_ignore_rules = parse_gitignore(gitignore) if gitignore else None
     dir_ = {root.name: _list_children(root=root, git_ignore_rules=git_ignore_rules)}
     v = _print_tree(dir_)
     return "\n".join(v)
@@ -110,3 +127,14 @@ def _add_line(rows: List[str]) -> List[str]:
             return rows
         rows[i] = "|" + v[1:]
     return rows
+
+
+def _execute_tree(root: Path, gitignore: str | Path) -> str:
+    args = ["--gitignore", str(gitignore)] if gitignore else []
+    try:
+        result = subprocess.run(["tree"] + args + [str(root)], capture_output=True, text=True, check=True)
+        if result.returncode != 0:
+            raise ValueError(f"tree exits with code {result.returncode}")
+        return result.stdout
+    except subprocess.CalledProcessError as e:
+        raise e
diff --git a/tests/metagpt/utils/test_tree.py b/tests/metagpt/utils/test_tree.py
index 34eae10cf..03a2a5606 100644
--- a/tests/metagpt/utils/test_tree.py
+++ b/tests/metagpt/utils/test_tree.py
@@ -2,7 +2,6 @@ from pathlib import Path
 from typing import List
 
 import pytest
-from gitignore_parser import parse_gitignore
 
 from metagpt.utils.tree import _print_tree, tree
 
@@ -15,8 +14,19 @@ from metagpt.utils.tree import _print_tree, tree
     ],
 )
 def test_tree(root: str, rules: str):
-    gitignore_rules = parse_gitignore(full_path=rules) if rules else None
-    v = tree(root=root, git_ignore_rules=gitignore_rules)
+    v = tree(root=root, gitignore=rules)
+    assert v
+
+
+@pytest.mark.parametrize(
+    ("root", "rules"),
+    [
+        (str(Path(__file__).parent / "../.."), None),
+        (str(Path(__file__).parent / "../.."), str(Path(__file__).parent / "../../../.gitignore")),
+    ],
+)
+def test_tree_command(root: str, rules: str):
+    v = tree(root=root, gitignore=rules, run_command=True)
     assert v