Add the --depth option

README.md

@@ -49,9 +49,12 @@ RP Tree also provides the following options:
 - `-h`, `--help` show a usage message
 - `-d`, `--dir-only` generates a directory-only tree diagram
 - `-o`, `--output-file` generates a full directory tree diagram and save it to a file in markdown format
+- `--depth N` limit the depth of the tree to N levels
 
 ## Release History
 
+- 1.0.0
+  - Add a `--depth` option to limit the depth of the tree
 - 0.1.1
   - Display the entries in alphabetical order
 - 0.1.0

pyproject.toml

@@ -0,0 +1,45 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "rptree"
+dynamic = ["version"]
+description = "Generate directory tree diagrams in the terminal"
+readme = "README.md"
+requires-python = ">=3.10"
+authors = [
+    {name = "Real Python", email = "[email protected]"},
+]
+maintainers = [
+    {name = "Leodanis Pozo Ramos", email = "[email protected]"},
+]
+license = {text = "MIT"}
+classifiers = [
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.7",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: Implementation :: CPython",
+]
+
+dependencies = []  # Add dependencies if required
+
+[project.scripts]
+rptree = "rptree.__main__:main"
+
+[tool.setuptools]
+packages = ["rptree"]
+include-package-data = true
+
+[tool.setuptools.dynamic]
+version = {attr = "rptree.__version__"}
+
+[tool.setuptools.exclude-package-data]
+"*" = ["tests/*", "tests/**/*"]

rptree/__init__.py

@@ -1,3 +1,3 @@
 """Top-level package for RP Tree."""
 
-__version__ = "0.1.1"
+__version__ = "1.0.0"

rptree/__main__.py

@@ -1,20 +1,25 @@
-"""This module provides the RP Tree CLI."""
+"""This module provides RP Tree entry point script."""
 
 import pathlib
 import sys
 
-from .cli import parse_cmd_line_arguments
+from . import cli
 from .rptree import DirectoryTree
 
 
-def main():
-    args = parse_cmd_line_arguments()
+def main() -> None:
+    args = cli.parse_cmd_line_arguments()
+
     root_dir = pathlib.Path(args.root_dir)
     if not root_dir.is_dir():
         print("The specified root directory doesn't exist")
-        sys.exit()
+        sys.exit(1)
+
     tree = DirectoryTree(
-        root_dir, dir_only=args.dir_only, output_file=args.output_file
+        root_dir=args.root_dir,
+        dir_only=args.dir_only,
+        output_file=args.output_file,
+        max_depth=args.depth,
     )
     tree.generate()
 

rptree/cli.py

@@ -6,33 +6,45 @@
 from . import __version__
 
 
-def parse_cmd_line_arguments():
+def parse_cmd_line_arguments() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         prog="tree",
-        description="RP Tree, a directory tree generator",
+        description="RP Tree, a directory tree diagram generator",
         epilog="Thanks for using RP Tree!",
     )
-    parser.version = f"RP Tree v{__version__}"
-    parser.add_argument("-v", "--version", action="version")
+    parser.add_argument(
+        "-v",
+        "--version",
+        action="version",
+        version=f"RP Tree v{__version__}",
+        help="displays the version of RP Tree",
+    )
     parser.add_argument(
         "root_dir",
         metavar="ROOT_DIR",
         nargs="?",
         default=".",
-        help="generate a full directory tree starting at ROOT_DIR",
+        help="generates a directory tree diagram starting at ROOT_DIR",
     )
     parser.add_argument(
         "-d",
         "--dir-only",
         action="store_true",
-        help="generate a directory-only tree",
+        help="generates a directory-only (no files) tree diagram",
     )
     parser.add_argument(
         "-o",
         "--output-file",
         metavar="OUTPUT_FILE",
         nargs="?",
         default=sys.stdout,
-        help="generate a full directory tree and save it to a file",
+        help="saves the directory tree diagram to OUTPUT_FILE in Markdown format",
+    )
+    parser.add_argument(
+        "--depth",
+        type=int,
+        metavar="N",
+        default=None,
+        help="limits the depth of the directory tree diagram to N levels",
     )
     return parser.parse_args()

rptree/rptree.py

@@ -4,6 +4,7 @@
 import pathlib
 import sys
 from collections import deque
+from typing import TextIO
 
 PIPE = "│"
 ELBOW = "└──"
@@ -13,39 +14,62 @@
 
 
 class DirectoryTree:
-    def __init__(self, root_dir, dir_only=False, output_file=sys.stdout):
+    def __init__(
+        self,
+        root_dir: str,
+        dir_only: bool = False,
+        output_file: str | TextIO = sys.stdout,
+        max_depth: int | None = None,
+    ) -> None:
         self._output_file = output_file
-        self._generator = _TreeGenerator(root_dir, dir_only)
+        self._generator = _TreeGenerator(root_dir, dir_only, max_depth)
 
-    def generate(self):
+    def generate(self) -> None:
         tree = self._generator.build_tree()
-        if self._output_file != sys.stdout:
+
+        if isinstance(self._output_file, str):
             # Wrap the tree in a markdown code block
             tree.appendleft("```")
             tree.append("```")
-            self._output_file = open(
-                self._output_file, mode="w", encoding="UTF-8"
-            )
-        with self._output_file as stream:
+
+            with open(self._output_file, mode="w", encoding="UTF-8") as stream:
+                for entry in tree:
+                    print(entry, file=stream)
+        else:
             for entry in tree:
-                print(entry, file=stream)
+                print(entry, file=self._output_file)
 
 
 class _TreeGenerator:
-    def __init__(self, root_dir, dir_only=False):
+    def __init__(
+        self,
+        root_dir: str,
+        dir_only: bool = False,
+        max_depth: int | None = None,
+    ) -> None:
         self._root_dir = pathlib.Path(root_dir)
         self._dir_only = dir_only
-        self._tree = deque()
+        self._max_depth = max_depth
+        self._tree: deque[str] = deque()
 
-    def build_tree(self):
+    def build_tree(self) -> deque[str]:
         self._tree_head()
-        self._tree_body(self._root_dir)
+        self._tree_body(self._root_dir, current_depth=0)
         return self._tree
 
-    def _tree_head(self):
+    def _tree_head(self) -> None:
         self._tree.append(f"{self._root_dir}{os.sep}")
 
-    def _tree_body(self, directory, prefix=""):
+    def _tree_body(
+        self,
+        directory: pathlib.Path,
+        prefix: str = "",
+        current_depth: int = 0,
+    ) -> None:
+        # Stop recursion if max_depth is reached
+        if self._max_depth is not None and current_depth >= self._max_depth:
+            return
+
         entries = self._prepare_entries(directory)
         last_index = len(entries) - 1
         for index, entry in enumerate(entries):
@@ -54,22 +78,31 @@ def _tree_body(self, directory, prefix=""):
                 if index == 0:
                     self._tree.append(prefix + PIPE)
                 self._add_directory(
-                    entry, index, last_index, prefix, connector
+                    entry,
+                    index,
+                    last_index,
+                    prefix,
+                    connector,
+                    current_depth,
                 )
             else:
                 self._add_file(entry, prefix, connector)
 
-    def _prepare_entries(self, directory):
-        entries = sorted(
-            directory.iterdir(), key=lambda entry: str(entry)
-        )
+    def _prepare_entries(self, directory: pathlib.Path) -> list[pathlib.Path]:
+        entries = sorted(directory.iterdir(), key=lambda entry: str(entry))
         if self._dir_only:
             return [entry for entry in entries if entry.is_dir()]
         return sorted(entries, key=lambda entry: entry.is_file())
 
     def _add_directory(
-        self, directory, index, last_index, prefix, connector
-    ):
+        self,
+        directory: pathlib.Path,
+        index: int,
+        last_index: int,
+        prefix: str,
+        connector: str,
+        current_depth: int,
+    ) -> None:
         self._tree.append(f"{prefix}{connector} {directory.name}{os.sep}")
         if index != last_index:
             prefix += PIPE_PREFIX
@@ -78,9 +111,10 @@ def _add_directory(
         self._tree_body(
             directory=directory,
             prefix=prefix,
+            current_depth=current_depth + 1,
         )
         if prefix := prefix.rstrip():
             self._tree.append(prefix)
 
-    def _add_file(self, file, prefix, connector):
+    def _add_file(self, file: pathlib.Path, prefix: str, connector: str) -> None:
         self._tree.append(f"{prefix}{connector} {file.name}")

tests/test_rptree.py

@@ -0,0 +1,67 @@
+from collections import deque
+from io import StringIO
+
+import pytest
+
+from rptree.rptree import DirectoryTree, _TreeGenerator
+
+
+@pytest.fixture
+def temp_dir(tmp_path):
+    """Create a temporary directory structure for testing."""
+    (tmp_path / "dir1").mkdir()
+    (tmp_path / "dir2").mkdir()
+    (tmp_path / "dir1" / "file1.txt").write_text("content")
+    (tmp_path / "dir2" / "file2.txt").write_text("content")
+    (tmp_path / "file3.txt").write_text("content")
+    return tmp_path
+
+
+def test_directory_tree_full(temp_dir):
+    """Test generating a full directory tree including files."""
+    output = StringIO()
+    tree = DirectoryTree(str(temp_dir), output_file=output)
+    tree.generate()
+    result = output.getvalue()
+    assert "dir1/" in result
+    assert "dir2/" in result
+    assert "file3.txt" in result
+    assert "file1.txt" in result
+    assert "file2.txt" in result
+
+
+def test_directory_tree_dir_only(temp_dir):
+    """Test generating a directory tree with only directories."""
+    output = StringIO()
+    tree = DirectoryTree(str(temp_dir), dir_only=True, output_file=output)
+    tree.generate()
+    result = output.getvalue()
+    assert "dir1/" in result
+    assert "dir2/" in result
+    assert "file3.txt" not in result
+    assert "file1.txt" not in result
+    assert "file2.txt" not in result
+
+
+def test_directory_tree_max_depth(temp_dir):
+    """Test generating a directory tree with max depth constraint."""
+    output = StringIO()
+    tree = DirectoryTree(str(temp_dir), output_file=output, max_depth=1)
+    tree.generate()
+    result = output.getvalue()
+    assert "dir1/" in result
+    assert "dir2/" in result
+    assert "file3.txt" in result
+    assert "file1.txt" not in result
+    assert "file2.txt" not in result
+
+
+def test_tree_generator(temp_dir):
+    """Test the _TreeGenerator class separately."""
+    generator = _TreeGenerator(str(temp_dir))
+    tree = generator.build_tree()
+    assert isinstance(tree, deque)
+    assert len(tree) > 0
+    assert any("dir1/" in line for line in tree)
+    assert any("dir2/" in line for line in tree)
+    assert any("file3.txt" in line for line in tree)