Add --levels/-L option

Signed-off-by: John Pennycook <[email protected]>

Author: Pennycook

codebasin/report.py

@@ -649,9 +649,11 @@ def insert(
     def _print(
         self,
         node: Node,
+        depth: int = 0,
         prefix: str = "",
         connector: str = "",
         fancy: bool = True,
+        levels: int = None,
     ):
         """
         Recursive helper function to print all nodes in a FileTree.
@@ -669,7 +671,14 @@ def _print(
 
         fancy: bool, default: True
             Whether to use fancy formatting (including colors).
+
+        levels: int, optional
+            The maximum number of levels to print.
         """
+        # Skip this node and its children if we have hit the maximum depth.
+        if levels and depth > levels:
+            return []
+
         if fancy:
             dash = "\u2500"
             cont = "\u251C"
@@ -722,23 +731,29 @@ def _print(
                 next_connector = cont
             lines += self._print(
                 node.children[name],
+                depth + 1,
                 next_prefix,
                 next_connector,
                 fancy,
+                levels,
             )
 
         return lines
 
-    def write_to(self, stream: TextIO):
+    def write_to(self, stream: TextIO, levels: int = None):
         """
         Write the FileTree to the specified stream.
 
         Parameters
         ----------
         stream: TextIO
             The text stream to write to.
+
+        levels: int, optional
+            The maximum number of levels to print.
+            If 0, print only the top-level summary.
         """
-        lines = self._print(self.root, fancy=stream.isatty())
+        lines = self._print(self.root, fancy=stream.isatty(), levels=levels)
         output = "\n".join(lines)
         if not stream.isatty():
             output = _strip_colors(output)
@@ -751,6 +766,7 @@ def files(
     *,
     stream: TextIO = sys.stdout,
     prune: bool = False,
+    levels: int = None,
 ):
     """
     Produce a file tree representing the code base.
@@ -815,4 +831,4 @@ def files(
     print(legend, file=stream)
 
     # Print the tree.
-    tree.write_to(stream)
+    tree.write_to(stream, levels=levels)

codebasin/tree.py

@@ -70,6 +70,17 @@ def _build_parser() -> argparse.ArgumentParser:
         help=_help_string(
             "Prune unused files from the tree.",
             is_long=True,
+        ),
+    )
+    parser.add_argument(
+        "-L",
+        "--levels",
+        dest="levels",
+        metavar="<level>",
+        type=int,
+        help=_help_string(
+            "Print only the specified number of levels.",
+            is_long=True,
             is_last=True,
         ),
     )
@@ -88,6 +99,10 @@ def _build_parser() -> argparse.ArgumentParser:
 
 
 def _tree(args: argparse.Namespace):
+    # Refuse to print a tree with no levels, consistent with tree utility.
+    if args.levels is not None and args.levels <= 0:
+        raise ValueError("Number of levels must be greater than 0.")
+
     # TODO: Refactor this to avoid duplication in __main__
     # Determine the root directory based on where codebasin is run.
     rootdir = os.path.abspath(os.getcwd())
@@ -140,7 +155,7 @@ def _tree(args: argparse.Namespace):
     )
 
     # Print the file tree.
-    report.files(codebase, state, prune=args.prune)
+    report.files(codebase, state, prune=args.prune, levels=args.levels)
     sys.exit(0)