Code clean up

README.md

@@ -4,13 +4,24 @@ RP Tree is a command-line tool to generate a directory tree diagram.
 
 ## Run the App
 
-Create a Python virtual environment and run the application:
+To run **RP Tree**, you need to download the source code. Then open a terminal or command-line window and run the following steps:
+
+1. Create and activate a Python virtual environment
 
 ```sh
-$ python3 tree.py /path/to/directory/
+$ cd rptree/
+$ python -m venv ./venv
+$ source venv/bin/activate
+(venv) $
 ```
 
-**Note:** The `-h` option provides help on how to use RP Tree.
+2. Run the application
+
+```sh
+(venv) $ python tree.py /path/to/directory/
+```
+
+**Note:** The `-h` or `--help` option provides help on how to use RP Tree.
 
 ## Current Features
 
@@ -23,6 +34,11 @@ RP Tree also provides the following options:
 - `-d`, `--dir-only` generates a directory-only tree
 - `-o`, `--output-file` generates a tree and save it to a file in markdown format
 
+## Release History
+
+- 0.1.0
+  - A work in progress
+
 ## About the Author
 
 Leodanis Pozo Ramos - Email: leodanis@realpython.com

rptree/cli.py

@@ -1,4 +1,4 @@
-"""This module provides RP Tree cli."""
+"""This module provides the RP Tree CLI."""
 
 import argparse
 import pathlib
@@ -11,44 +11,42 @@
 def main():
     args = 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()
-
-    tree = DirectoryTree(root_dir, args.output_file, args.dir_only)
+    tree = DirectoryTree(
+        root_dir, dir_only=args.dir_only, output_file=args.output_file
+    )
     tree.generate()
 
 
 def parse_cmd_line_arguments():
     parser = argparse.ArgumentParser(
-        prog="RP Tree",
-        description="Directory Tree Generator",
-        epilog="Thanks for Using RP Tree!",
+        prog="tree",
+        description="RP Tree, a directory tree generator",
+        epilog="Thanks for using RP Tree!",
     )
-    parser.version = f"{parser.prog} {__version__}"
+    parser.version = f"RP Tree v{__version__}"
+    parser.add_argument("-v", "--version", action="version")
     parser.add_argument(
         "root_dir",
         metavar="ROOT_DIR",
         nargs="?",
         default=".",
-        type=str,
-        help="Generate a Full and Recursive Directory Tree",
+        help="Generate a full directory tree starting at ROOT_DIR",
     )
     parser.add_argument(
         "-d",
         "--dir-only",
         action="store_true",
-        help="Generate a Directory-Only Tree",
+        help="Generate a directory-only tree",
     )
     parser.add_argument(
         "-o",
         "--output-file",
         metavar="OUTPUT_FILE",
         nargs="?",
         default=sys.stdout,
-        type=str,
-        help="Generate a Full Directory Tree and Save it to a File",
+        help="Generate a full directory tree and save it to a file",
     )
-    parser.add_argument("-v", "--version", action="version")
     return parser.parse_args()

rptree/rptree.py

@@ -4,75 +4,82 @@
 import pathlib
 import sys
 
+from collections import deque
+
 PIPE = "│"
 ELBOW = "└──"
 TEE = "├──"
 PIPE_PREFIX = "│   "
+SPACE_PREFIX = "    "
 
 
 class DirectoryTree:
-    def __init__(self, root_dir, output_file=sys.stdout, dir_only=False):
+    def __init__(self, root_dir, dir_only=False, output_file=sys.stdout):
         self._output_file = output_file
         self._generator = _TreeGenerator(root_dir, dir_only)
 
     def generate(self):
-        tree = self._generator.generate_tree()
+        tree = self._generator.build_tree()
         if self._output_file != sys.stdout:
+            # 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:
-            for item in tree:
-                print(item, file=stream)
+            for entry in tree:
+                print(entry, file=stream)
 
 
 class _TreeGenerator:
-    def __init__(self, root, dir_only=False):
-        self._root = pathlib.Path(root)
-        self.dir_only = dir_only
-        self._tree = []
+    def __init__(self, root_dir, dir_only=False):
+        self._root_dir = pathlib.Path(root_dir)
+        self._dir_only = dir_only
+        self._tree = deque()
 
-    def generate_tree(self):
+    def build_tree(self):
         self._tree_head()
-        self._tree_tail(self._root)
+        self._tree_body(self._root_dir)
         return self._tree
 
     def _tree_head(self):
-        self._tree.append(f"{self._root}{os.sep}")
+        self._tree.append(f"{self._root_dir}{os.sep}")
         self._tree.append(PIPE)
 
-    def _tree_tail(self, dir, prefix=""):
-        entries = self._prepare_entries(dir)
+    def _tree_body(self, directory, prefix=""):
+        entries = self._prepare_entries(directory)
         entries_count = len(entries)
         for index, entry in enumerate(entries):
             connector = ELBOW if index == entries_count - 1 else TEE
-            path = dir.joinpath(entry)
-            if path.is_dir():
+            if entry.is_dir():
                 self._add_directory(
                     entry, index, entries_count, prefix, connector
                 )
             else:
-                self._add_file(prefix, connector, path)
+                self._add_file(entry, prefix, connector)
 
-    def _prepare_entries(self, dir):
-        entries = dir.iterdir()
-        if self.dir_only:
-            entries = [e for e in entries if dir.joinpath(e).is_dir()]
+    def _prepare_entries(self, directory):
+        entries = directory.iterdir()
+        if self._dir_only:
+            entries = [entry for entry in entries if entry.is_dir()]
             return entries
-        entries = sorted(entries, key=lambda e: dir.joinpath(e).is_file())
+        entries = sorted(entries, key=lambda entry: entry.is_file())
         return entries
 
-    def _add_directory(self, entry, index, entries_count, prefix, connector):
-        self._tree.append(f"{prefix}{connector}{entry.name}{os.sep}")
+    def _add_directory(
+        self, directory, index, entries_count, prefix, connector
+    ):
+        self._tree.append(f"{prefix}{connector} {directory.name}{os.sep}")
         if index != entries_count - 1:
             prefix += PIPE_PREFIX
-        self._tree_tail(
-            dir=entry,
+        else:
+            prefix += SPACE_PREFIX
+        self._tree_body(
+            directory=directory,
             prefix=prefix,
         )
-        self._tree.append(prefix)
+        self._tree.append(prefix.rstrip())
 
-    def _add_file(self, prefix, connector, path):
-        self._tree.append(f"{prefix}{connector} {path.name}")
+    def _add_file(self, file, prefix, connector):
+        self._tree.append(f"{prefix}{connector} {file.name}")

tree.py

@@ -4,6 +4,5 @@
 
 from rptree.cli import main
 
-
 if __name__ == "__main__":
     main()