diff --git a/MANIFEST.in b/MANIFEST.in deleted file mode 100644 index 1aba38f..0000000 --- a/MANIFEST.in +++ /dev/null @@ -1 +0,0 @@ -include LICENSE diff --git a/README.md b/README.md index f647053..ec87833 100644 --- a/README.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..e104757 --- /dev/null +++ b/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 = "info@realpython.com"}, +] +maintainers = [ + {name = "Leodanis Pozo Ramos", email = "leodanis@realpython.com"}, +] +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/**/*"] \ No newline at end of file diff --git a/rptree/__init__.py b/rptree/__init__.py index e8d5a69..13c8700 100644 --- a/rptree/__init__.py +++ b/rptree/__init__.py @@ -1,3 +1,3 @@ """Top-level package for RP Tree.""" -__version__ = "0.1.1" +__version__ = "1.0.0" diff --git a/rptree/__main__.py b/rptree/__main__.py index 0613440..8726346 100644 --- a/rptree/__main__.py +++ b/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() diff --git a/rptree/cli.py b/rptree/cli.py index c12eb43..74615b0 100644 --- a/rptree/cli.py +++ b/rptree/cli.py @@ -6,26 +6,31 @@ 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", @@ -33,6 +38,13 @@ def parse_cmd_line_arguments(): 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() diff --git a/rptree/rptree.py b/rptree/rptree.py index 2581c0a..2dd01af 100644 --- a/rptree/rptree.py +++ b/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}") diff --git a/setup.py b/setup.py deleted file mode 100644 index 3ddee53..0000000 --- a/setup.py +++ /dev/null @@ -1,38 +0,0 @@ -import pathlib -from setuptools import setup - -from rptree import __version__ - -HERE = pathlib.Path().cwd() -DESCRIPTION = HERE.joinpath("README.md").read_text() -VERSION = __version__ - - -setup( - name="rptree", - version=VERSION, - description="Generate directory tree diagrams for Real Python articles", - long_description=DESCRIPTION, - long_description_content_type="text/markdown", - url="https://github.com/realpython/rptree", - author="Real Python", - author_email="info@realpython.com", - maintainer="Leodanis Pozo Ramos", - maintainer_email="leodanis@realpython.com", - license="MIT", - classifiers=[ - "License :: OSI Approved :: MIT License", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.7", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", - "Programming Language :: Python :: Implementation :: CPython", - ], - packages=["rptree"], - include_package_data=True, - entry_points={ - "console_scripts": [ - "rptree=rptree.__main__:main", - ] - }, -) diff --git a/tests/test_rptree.py b/tests/test_rptree.py new file mode 100644 index 0000000..69b135f --- /dev/null +++ b/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)