diff --git a/scripts/bump_dataset_version.py b/scripts/bump_dataset_version.py index 0b61ce3..1a2732b 100644 --- a/scripts/bump_dataset_version.py +++ b/scripts/bump_dataset_version.py @@ -20,11 +20,33 @@ import datetime as dt import os import re import subprocess +import sys from pathlib import Path GEN_RE = re.compile(r"^gen(\d+)$") SCHEMA_RE = re.compile(r"^schema(\d+)$") +# One color per tree level in `dwarf status` output, so the eye can jump +# straight to e.g. "all the schema rows" or "all the totals". +_LEVEL_COLORS = { + "kind": "\033[1;36m", # bold cyan — top-level kind/ header + "gen": "\033[1;33m", # bold yellow — genN row + kind total + "bucket": "\033[34m", # blue — raw/processed subtotals + "schema": "\033[32m", # green — schemaN rows + "root": "\033[1;35m", # bold magenta — derived/, pools/, grand total +} +_RESET = "\033[0m" + + +def _use_color() -> bool: + return sys.stdout.isatty() and os.environ.get("NO_COLOR") is None + + +def _colorize(text: str, level: str) -> str: + if not _use_color(): + return text + return f"{_LEVEL_COLORS[level]}{text}{_RESET}" + def _find_component_idx(abs_path: Path, pattern: re.Pattern) -> int | None: """Return the index of the first path component matching *pattern*, or None.""" @@ -147,11 +169,12 @@ def _human_size(n: int) -> str: _ROW_WIDTH = 48 -def _row(label: str, size_bytes: int, indent: int = 0) -> str: +def _row(label: str, size_bytes: int, indent: int = 0, level: str | None = None) -> str: text = " " * indent + label size_str = _human_size(size_bytes) pad = max(1, _ROW_WIDTH - len(text) - len(size_str)) - return text + " " * pad + size_str + row = text + " " * pad + size_str + return _colorize(row, level) if level else row def print_status(root: Path) -> None: @@ -167,7 +190,7 @@ def print_status(root: Path) -> None: for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir()) if m ) - print(f"{kind}/") + print(_colorize(f"{kind}/", "kind")) kind_total = 0 for gen in gens: gen_tag = f"gen{gen}" @@ -187,25 +210,25 @@ def print_status(root: Path) -> None: gen_total = raw_size + processed_size kind_total += gen_total - print(_row(gen_tag, gen_total, indent=1)) - print(_row("raw", raw_size, indent=2)) - print(_row("processed", processed_size, indent=2)) + print(_row(gen_tag, gen_total, indent=1, level="gen")) + print(_row("raw", raw_size, indent=2, level="bucket")) + print(_row("processed", processed_size, indent=2, level="bucket")) if schemas: for s in schemas: - print(_row(f"schema{s}", schema_sizes[s], indent=3)) + print(_row(f"schema{s}", schema_sizes[s], indent=3, level="schema")) else: - print(" (none)") - print(_row(f"{kind} total", kind_total, indent=1)) + print(_colorize(" (none)", "schema")) + print(_row(f"{kind} total", kind_total, indent=1, level="gen")) print() grand_total += kind_total derived_size = _du(root / "derived") pools_size = _du(root / "pools") grand_total += derived_size + pools_size - print(_row("derived/", derived_size)) - print(_row("pools/", pools_size)) + print(_row("derived/", derived_size, level="root")) + print(_row("pools/", pools_size, level="root")) print("-" * _ROW_WIDTH) - print(_row("grand total", grand_total)) + print(_row("grand total", grand_total, level="root")) # ---------------------------------------------------------------------------