make tree recursive

2 files changed, 79 insertions, 6 deletions

diff --git a/owners/_colors.py b/owners/_colors.py
index 43e034b..c40e4a4 100644
--- a/owners/_colors.py
+++ b/owners/_colors.py
@@ -6,6 +6,7 @@ from typing import Any
 
 RED = '\033[31m'
 GREEN = '\033[32m'
+YELLOW = '\033[33m'
 BLUE = '\033[94m'
 MAGENTA = '\033[35m'
 END = '\033[0m'
@@ -34,3 +35,8 @@ class Colors:
         if self.disabled:
             return str(text)
         return f'{MAGENTA}{text}{END}'
+
+    def yellow(self, text: Any) -> str:
+        if self.disabled:
+            return str(text)
+        return f'{YELLOW}{text}{END}'
diff --git a/owners/commands/_tree.py b/owners/commands/_tree.py
index 2963b51..c8fc308 100644
--- a/owners/commands/_tree.py
+++ b/owners/commands/_tree.py
@@ -1,10 +1,50 @@
 from __future__ import annotations
 
 from argparse import ArgumentParser
+from dataclasses import dataclass
+from functools import cached_property
 from pathlib import Path
+from typing import Callable
 
 from ._base import Command
 
+from enum import Enum
+
+
+class Cov(Enum):
+    FULL = 'full'  # the dir is 100% owned
+    MOST = 'most'  # not the dir itself but everything inside is owned
+    PART = 'part'  # some subdirs in the dir are owned
+    NONE = 'none'  # nothing in the dir is owned
+
+
+# Directories to skip by default
+IGNORE = frozenset({'.git', '__pycache__'})
+
+
+@dataclass
+class PathInfo:
+    # path to the analyzed dir
+    path: Path
+
+    # information about all subdirs, None if not analyzed (because the root is owned)
+    children: list[PathInfo] | None
+
+    # the owner team, empty string for not directly owned dirs
+    owner: str
+
+    @cached_property
+    def cov(self) -> Cov:
+        if self.children is None:
+            return Cov.FULL
+        if not self.children:
+            return Cov.NONE
+        if all(p.cov in (Cov.FULL, Cov.MOST) for p in self.children):
+            return Cov.MOST
+        if all(p.cov == Cov.NONE for p in self.children):
+            return Cov.NONE
+        return Cov.PART
+
 
 class Tree(Command):
     """Show file tree with ownership info.
@@ -15,16 +55,43 @@ class Tree(Command):
         Command.init_parser(parser)
 
     def run(self) -> int:
-        self._inspect(self.code_owners.root)
+        path_info = self._inspect(self.code_owners.root)
+        self._show_tree(path_info, 0)
         return 0
 
-    def _inspect(self, root: Path) -> None:
+    def _inspect(self, root: Path) -> PathInfo:
+        rule = self.code_owners.find_rule(root)
+        if rule is not None:
+            return PathInfo(root, children=None, owner=rule.owners[0])
+
+        path_infos: list[PathInfo] = []
         for path in sorted(root.iterdir()):
             if not path.is_dir():
                 continue
-            rule = self.code_owners.find_rule(path)
-            if rule is not None:
-                self.print(self.colors.green(path.name))
+            if path.name in IGNORE:
                 continue
+            path_infos.append(self._inspect(path))
+        return PathInfo(root, children=path_infos, owner='')
+
+    def _show_tree(self, info: PathInfo, level: int) -> None:
+        if info.cov == Cov.FULL:
+            self._show_path(self.colors.green, info, level)
+            return
+        if info.cov == Cov.MOST:
+            self._show_path(self.colors.blue, info, level)
+            return
+        if info.cov == Cov.NONE:
+            self._show_path(self.colors.red, info, level)
+            return
+        if info.cov == Cov.PART:
+            assert info.children
+            self._show_path(self.colors.yellow, info, level)
+            for child in info.children:
+                self._show_tree(child, level + 1)
+            return
+        raise RuntimeError('unreachable')
 
-            self.print(self.colors.red(path.name))
+    def _show_path(self, color: Callable, info: PathInfo, level: int) -> None:
+        width = level * 2 + len(info.path.name) + 1
+        ident = 60 - width
+        self.print(' |' * level, color(info.path.name), ' ' * ident + info.owner)