Create _tree_utils.py

Modules/_tree_utils.py

@@ -0,0 +1,265 @@
+"""
+Shared tree rendering utilities for filesystem-like output.
+
+Provides functions to build and render tree structures with line connectors
+(├──, └──, │) for visual hierarchy display.
+"""
+
+from __future__ import annotations
+
+import os
+from datetime import datetime
+from typing import Callable, Optional
+
+
+def build_tree(entries: list[tuple[str, dict]]) -> dict:
+    """
+    Build a nested tree structure from flat path entries.
+    
+    Args:
+        entries: List of (path, metadata) tuples where path uses forward slashes.
+                 Paths ending with '/' are treated as directories.
+    
+    Returns:
+        Nested dict with "__files__" key for files at each level.
+    
+    Example:
+        entries = [
+            ("scripts/utils.py", {"size": 1234}),
+            ("docs/readme.md", {"size": 500}),
+        ]
+        tree = build_tree(entries)
+    """
+    root: dict = {"__files__": []}
+    
+    for path, metadata in entries:
+        parts = path.rstrip("/").split("/")
+        is_dir = path.endswith("/")
+        
+        node = root
+        for i, part in enumerate(parts[:-1]):
+            if part not in node:
+                node[part] = {"__files__": []}
+            node = node[part]
+        
+        final = parts[-1]
+        if is_dir:
+            # Ensure directory exists
+            if final not in node:
+                node[final] = {"__files__": []}
+            # Store directory metadata if provided
+            if metadata:
+                node[final]["__meta__"] = metadata
+        else:
+            # Add file
+            node["__files__"].append((final, metadata))
+    
+    return root
+
+
+def render_tree(
+    node: dict,
+    prefix: str = "",
+    format_entry: Optional[Callable[[str, dict, bool], str]] = None,
+) -> list[str]:
+    """
+    Render a tree with line connectors.
+    
+    Args:
+        node: Nested dict from build_tree()
+        prefix: Current line prefix for indentation
+        format_entry: Optional callback to format each entry.
+                      Signature: (name, metadata, is_dir) -> str
+                      If None, uses default formatting.
+    
+    Returns:
+        List of formatted lines.
+    """
+    result = []
+    
+    # Default formatter
+    def default_format(name: str, meta: dict, is_dir: bool) -> str:
+        if is_dir:
+            return f"{name}/"
+        size = meta.get("size")
+        if size is not None:
+            return f"{name}  ({_fmt_size(size)})"
+        return name
+    
+    fmt = format_entry or default_format
+    
+    # Collect entries: subdirs first, then files
+    entries = []
+    subdirs = sorted(k for k in node.keys() if k not in ("__files__", "__meta__"))
+    files_here = sorted(node.get("__files__", []), key=lambda x: x[0])
+    
+    for dirname in subdirs:
+        dir_meta = node[dirname].get("__meta__", {})
+        entries.append(("dir", dirname, node[dirname], dir_meta))
+    for fname, fmeta in files_here:
+        entries.append(("file", fname, None, fmeta))
+    
+    for i, entry in enumerate(entries):
+        is_last = (i == len(entries) - 1)
+        connector = "└── " if is_last else "├── "
+        child_prefix = prefix + ("    " if is_last else "│   ")
+        
+        etype, name, subtree, meta = entry
+        
+        if etype == "dir":
+            result.append(f"{prefix}{connector}{fmt(name, meta, True)}")
+            result.extend(render_tree(subtree, child_prefix, format_entry))
+        else:
+            result.append(f"{prefix}{connector}{fmt(name, meta, False)}")
+    
+    return result
+
+
+def _fmt_size(num_bytes: int) -> str:
+    """Format byte size as human-readable string."""
+    units = ["B", "KB", "MB", "GB"]
+    size = float(num_bytes)
+    for unit in units:
+        if size < 1024.0:
+            return f"{size:.1f} {unit}"
+        size /= 1024.0
+    return f"{size:.1f} TB"
+
+
+def walk_and_build_tree(
+    abs_path: str,
+    *,
+    show_hidden: bool = False,
+    recursive: bool = False,
+    max_entries: int = 100,
+) -> tuple[dict, int, bool]:
+    """
+    Walk a directory and build a tree structure.
+    
+    Args:
+        abs_path: Absolute path to directory
+        show_hidden: Include hidden files/dirs (starting with '.')
+        recursive: Recurse into subdirectories
+        max_entries: Maximum entries before truncation
+    
+    Returns:
+        (tree, total_entries, truncated)
+    """
+    entries: list[tuple[str, dict]] = []
+    total = 0
+    truncated = False
+    
+    for root, dirs, files in os.walk(abs_path):
+        # Filter hidden
+        if not show_hidden:
+            dirs[:] = [d for d in dirs if not d.startswith('.')]
+            files = [f for f in files if not f.startswith('.')]
+        
+        dirs.sort()
+        files.sort()
+        
+        # Compute relative path from the listing root
+        try:
+            rel_root = os.path.relpath(root, abs_path)
+        except Exception:
+            rel_root = ""
+        prefix = "" if rel_root == "." else rel_root.replace("\\", "/") + "/"
+        
+        # Add directories (with trailing slash to indicate dir)
+        for d in dirs:
+            p = os.path.join(root, d)
+            try:
+                mtime = datetime.fromtimestamp(os.path.getmtime(p)).strftime("%Y-%m-%d %H:%M")
+            except Exception:
+                mtime = "?"
+            entries.append((f"{prefix}{d}/", {"mtime": mtime}))
+            total += 1
+            if total >= max_entries:
+                truncated = True
+                break
+        
+        if truncated:
+            break
+        
+        # Add files
+        for f in files:
+            p = os.path.join(root, f)
+            try:
+                size = os.path.getsize(p)
+                mtime = datetime.fromtimestamp(os.path.getmtime(p)).strftime("%Y-%m-%d %H:%M")
+            except Exception:
+                size, mtime = 0, "?"
+            entries.append((f"{prefix}{f}", {"size": size, "mtime": mtime}))
+            total += 1
+            if total >= max_entries:
+                truncated = True
+                break
+        
+        if truncated:
+            break
+        
+        if not recursive:
+            break
+    
+    return build_tree(entries), total, truncated
+
+
+def format_dir_listing(
+    abs_path: str,
+    display_path: str,
+    *,
+    show_hidden: bool = False,
+    recursive: bool = False,
+    max_entries: int = 100,
+    fmt_size_fn: Optional[Callable[[int], str]] = None,
+) -> str:
+    """
+    Format a directory listing as a visual tree.
+    
+    Args:
+        abs_path: Absolute path to directory
+        display_path: User-friendly path to show in header
+        show_hidden: Include hidden files/dirs
+        recursive: Recurse into subdirectories
+        max_entries: Maximum entries before truncation
+        fmt_size_fn: Optional custom size formatter (defaults to _fmt_size)
+    
+    Returns:
+        Formatted string with tree output.
+    """
+    fmt_size = fmt_size_fn or _fmt_size
+    
+    tree, total, truncated = walk_and_build_tree(
+        abs_path,
+        show_hidden=show_hidden,
+        recursive=recursive,
+        max_entries=max_entries,
+    )
+    
+    # Formatter with size + date
+    def format_entry(name: str, meta: dict, is_dir: bool) -> str:
+        mtime = meta.get("mtime", "")
+        if is_dir:
+            return f"{name}/  ({mtime})"
+        size = meta.get("size", 0)
+        return f"{name}  ({fmt_size(size)}, {mtime})"
+    
+    tree_lines = render_tree(tree, "    ", format_entry)
+    
+    header = f"Listing of {display_path}\nRoot: /\nEntries: {total}"
+    if truncated:
+        header += f"\n… Truncated at {max_entries} entries."
+    
+    lines = [header, "", "└── /"]
+    lines.extend(tree_lines)
+    
+    return "\n".join(lines).strip()
+
+
+__all__ = [
+    "build_tree",
+    "render_tree",
+    "_fmt_size",
+    "walk_and_build_tree",
+    "format_dir_listing",
+]