Skip to content

Add Depth Limitation Flag to Directory Tree Generator #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Add Depth Limitation Flag to Directory Tree Generator #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
cc3ce11
Update rptree.py
sm18lr88 Jun 3, 2024
21cf1b8
Update cli.py
sm18lr88 Jun 3, 2024
9b2c54b
Update __main__.py
sm18lr88 Jun 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion rptree/__main__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
8 changes: 8 additions & 0 deletions rptree/cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
18 changes: 11 additions & 7 deletions rptree/rptree.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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):
Expand All @@ -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)
Expand All @@ -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:
Expand All @@ -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)
Expand Down