""" 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", ]