From cc3ce11f437a84b5b49d21b268d7551375e8dcd1 Mon Sep 17 00:00:00 2001
From: sm18lr88 <64564447+sm18lr88@users.noreply.github.com>
Date: Mon, 3 Jun 2024 19:49:16 -0400
Subject: [PATCH 1/3] Update rptree.py

---
 rptree/rptree.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/rptree/rptree.py b/rptree/rptree.py
index 2581c0a..ae73d90 100644
--- a/rptree/rptree.py
+++ b/rptree/rptree.py
@@ -13,9 +13,9 @@
 
 
 class DirectoryTree:
-    def __init__(self, root_dir, dir_only=False, output_file=sys.stdout):
+    def __init__(self, root_dir, dir_only=False, output_file=sys.stdout, max_depth=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):
         tree = self._generator.build_tree()
@@ -32,20 +32,23 @@ def generate(self):
 
 
 class _TreeGenerator:
-    def __init__(self, root_dir, dir_only=False):
+    def __init__(self, root_dir, dir_only=False, max_depth=None):
         self._root_dir = pathlib.Path(root_dir)
         self._dir_only = dir_only
         self._tree = deque()
+        self._max_depth = max_depth
 
     def build_tree(self):
         self._tree_head()
-        self._tree_body(self._root_dir)
+        self._tree_body(self._root_dir, depth=0)
         return self._tree
 
     def _tree_head(self):
         self._tree.append(f"{self._root_dir}{os.sep}")
 
-    def _tree_body(self, directory, prefix=""):
+    def _tree_body(self, directory, prefix="", depth=0):
+        if self._max_depth is not None and depth >= self._max_depth:
+            return
         entries = self._prepare_entries(directory)
         last_index = len(entries) - 1
         for index, entry in enumerate(entries):
@@ -54,7 +57,7 @@ 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, depth
                 )
             else:
                 self._add_file(entry, prefix, connector)
@@ -68,7 +71,7 @@ def _prepare_entries(self, directory):
         return sorted(entries, key=lambda entry: entry.is_file())
 
     def _add_directory(
-        self, directory, index, last_index, prefix, connector
+        self, directory, index, last_index, prefix, connector, depth
     ):
         self._tree.append(f"{prefix}{connector} {directory.name}{os.sep}")
         if index != last_index:
@@ -78,6 +81,7 @@ def _add_directory(
         self._tree_body(
             directory=directory,
             prefix=prefix,
+            depth=depth + 1,
         )
         if prefix := prefix.rstrip():
             self._tree.append(prefix)

From 21cf1b879f25fc4bb28c39f31340c4214a4beb20 Mon Sep 17 00:00:00 2001
From: sm18lr88 <64564447+sm18lr88@users.noreply.github.com>
Date: Mon, 3 Jun 2024 19:49:40 -0400
Subject: [PATCH 2/3] Update cli.py

adding the listing depth flag
---
 rptree/cli.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/rptree/cli.py b/rptree/cli.py
index c12eb43..210dcd9 100644
--- a/rptree/cli.py
+++ b/rptree/cli.py
@@ -35,4 +35,12 @@ def parse_cmd_line_arguments():
         default=sys.stdout,
         help="generate a full directory tree and save it to a file",
     )
+    parser.add_argument(
+        "-l",
+        "--level",
+        type=int,
+        metavar="LEVEL",
+        nargs="?",
+        help="limit the depth of the directory tree",
+    )
     return parser.parse_args()

From 9b2c54b90259ebcf1939f7bdd77447f752713eb1 Mon Sep 17 00:00:00 2001
From: sm18lr88 <64564447+sm18lr88@users.noreply.github.com>
Date: Mon, 3 Jun 2024 19:49:54 -0400
Subject: [PATCH 3/3] Update __main__.py

---
 rptree/__main__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/rptree/__main__.py b/rptree/__main__.py
index 0613440..f727182 100644
--- a/rptree/__main__.py
+++ b/rptree/__main__.py
@@ -14,7 +14,7 @@ def main():
         print("The specified root directory doesn't exist")
         sys.exit()
     tree = DirectoryTree(
-        root_dir, dir_only=args.dir_only, output_file=args.output_file
+        root_dir, dir_only=args.dir_only, output_file=args.output_file, max_depth=args.level
     )
     tree.generate()