diskuse.py

#!/usr/bin/env python3
"""
diskuse -- Analyze disk usage. Like du/ncdu, zero deps, cross-platform.

Find what's eating your disk. Top files, size by extension, directory breakdown.
Works on Windows, macOS, Linux without installing anything.

Usage:
    py diskuse.py .                            # Directory size summary
    py diskuse.py . --top 20                   # Top 20 largest files
    py diskuse.py . --ext                      # Size by file extension
    py diskuse.py . --dirs                     # Size by subdirectory
    py diskuse.py . --bigger 10M               # Files bigger than 10MB
    py diskuse.py . --smaller 1K               # Files smaller than 1KB
    py diskuse.py . --older 30                 # Files not modified in 30 days
    py diskuse.py . --ext --top 10             # Top 10 extensions by size
    py diskuse.py . --format json              # JSON output
"""

import argparse
import json
import os
import sys
import time
from collections import defaultdict
from datetime import datetime, timedelta
from pathlib import Path

# Colors
RESET = "\033[0m"
BOLD = "\033[1m"
DIM = "\033[2m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
CYAN = "\033[36m"
RED = "\033[31m"
MAGENTA = "\033[35m"

SKIP_DIRS = {".git", "node_modules", "__pycache__", ".venv", "venv"}


def color_supported() -> bool:
    if os.environ.get("NO_COLOR"):
        return False
    if sys.platform == "win32":
        return bool(os.environ.get("TERM") or os.environ.get("WT_SESSION"))
    return hasattr(sys.stdout, "isatty") and sys.stdout.isatty()


USE_COLOR = color_supported()


def c(code: str, text: str) -> str:
    return f"{code}{text}{RESET}" if USE_COLOR else text


def fmt_size(n: int) -> str:
    """Format bytes into human-readable string."""
    if n < 1024:
        return f"{n} B"
    elif n < 1024 ** 2:
        return f"{n / 1024:.1f} KB"
    elif n < 1024 ** 3:
        return f"{n / 1024**2:.1f} MB"
    else:
        return f"{n / 1024**3:.2f} GB"


def parse_size(s: str) -> int:
    """Parse human size string to bytes (e.g., '10M', '1.5G', '500K')."""
    s = s.strip().upper()
    multipliers = {"B": 1, "K": 1024, "KB": 1024, "M": 1024**2, "MB": 1024**2,
                   "G": 1024**3, "GB": 1024**3}
    for suffix, mult in sorted(multipliers.items(), key=lambda x: -len(x[0])):
        if s.endswith(suffix):
            try:
                return int(float(s[:-len(suffix)]) * mult)
            except ValueError:
                break
    try:
        return int(s)
    except ValueError:
        return 0


def bar(value: int, max_value: int, width: int = 25) -> str:
    """ASCII bar chart."""
    if max_value <= 0:
        return " " * width
    filled = int(width * value / max_value)
    return "#" * filled + " " * (width - filled)


def scan_directory(path: str, skip_hidden: bool = False) -> list[dict]:
    """Scan directory and collect file info."""
    files = []
    for root, dirs, filenames in os.walk(path):
        # Skip hidden/ignored directories
        dirs[:] = [d for d in dirs if d not in SKIP_DIRS]
        if skip_hidden:
            dirs[:] = [d for d in dirs if not d.startswith(".")]

        for fname in filenames:
            if skip_hidden and fname.startswith("."):
                continue
            fpath = os.path.join(root, fname)
            try:
                stat = os.stat(fpath)
                files.append({
                    "path": fpath,
                    "name": fname,
                    "size": stat.st_size,
                    "mtime": stat.st_mtime,
                    "ext": Path(fname).suffix.lower() or "(none)",
                    "dir": os.path.relpath(root, path),
                })
            except (OSError, PermissionError):
                pass
    return files


def cmd_summary(files: list[dict], path: str):
    """Show overall summary."""
    total_size = sum(f["size"] for f in files)
    file_count = len(files)
    dirs = set(f["dir"] for f in files)

    print(f"\n  {c(BOLD, 'Disk Usage Summary')}")
    print(f"  Path:        {c(CYAN, path)}")
    print(f"  Total size:  {c(BOLD, fmt_size(total_size))}")
    print(f"  Files:       {file_count:,}")
    print(f"  Directories: {len(dirs):,}")
    if files:
        largest = max(files, key=lambda f: f["size"])
        print(f"  Largest:     {fmt_size(largest['size'])}  {c(DIM, largest['name'])}")
    print()


def cmd_top(files: list[dict], n: int):
    """Show top N largest files."""
    sorted_files = sorted(files, key=lambda f: f["size"], reverse=True)[:n]
    if not sorted_files:
        print(c(DIM, "  No files found."))
        return

    max_size = sorted_files[0]["size"] if sorted_files else 1
    print(f"\n  {c(BOLD, f'Top {len(sorted_files)} Largest Files')}\n")
    for f in sorted_files:
        b = bar(f["size"], max_size)
        rel = os.path.relpath(f["path"])
        print(f"  {fmt_size(f['size']):>10}  {b}  {rel}")
    total = sum(f["size"] for f in sorted_files)
    print(f"\n  {c(DIM, f'Total: {fmt_size(total)}')}\n")


def cmd_ext(files: list[dict], n: int | None = None):
    """Show size breakdown by file extension."""
    by_ext: dict[str, dict] = defaultdict(lambda: {"size": 0, "count": 0})
    for f in files:
        by_ext[f["ext"]]["size"] += f["size"]
        by_ext[f["ext"]]["count"] += 1

    sorted_exts = sorted(by_ext.items(), key=lambda x: x[1]["size"], reverse=True)
    if n:
        sorted_exts = sorted_exts[:n]

    if not sorted_exts:
        return

    max_size = sorted_exts[0][1]["size"]
    total_size = sum(f["size"] for f in files)

    print(f"\n  {c(BOLD, 'Size by Extension')}\n")
    for ext, data in sorted_exts:
        pct = data["size"] / total_size * 100 if total_size else 0
        b = bar(data["size"], max_size)
        print(f"  {ext:>8}  {fmt_size(data['size']):>10}  ({pct:5.1f}%)  "
              f"{data['count']:>5} files  {b}")
    print()


def cmd_dirs(files: list[dict], n: int | None = None):
    """Show size breakdown by subdirectory."""
    by_dir: dict[str, dict] = defaultdict(lambda: {"size": 0, "count": 0})
    for f in files:
        # Use top-level directory
        top_dir = f["dir"].split(os.sep)[0] if f["dir"] != "." else "."
        by_dir[top_dir]["size"] += f["size"]
        by_dir[top_dir]["count"] += 1

    sorted_dirs = sorted(by_dir.items(), key=lambda x: x[1]["size"], reverse=True)
    if n:
        sorted_dirs = sorted_dirs[:n]

    if not sorted_dirs:
        return

    max_size = sorted_dirs[0][1]["size"]
    total_size = sum(f["size"] for f in files)

    print(f"\n  {c(BOLD, 'Size by Directory')}\n")
    for dirname, data in sorted_dirs:
        pct = data["size"] / total_size * 100 if total_size else 0
        b = bar(data["size"], max_size)
        print(f"  {dirname:<25}  {fmt_size(data['size']):>10}  ({pct:5.1f}%)  "
              f"{data['count']:>5} files  {b}")
    print()


def cmd_bigger(files: list[dict], threshold: int):
    """Show files bigger than threshold."""
    big = [f for f in files if f["size"] > threshold]
    big.sort(key=lambda f: f["size"], reverse=True)

    if not big:
        print(c(DIM, f"  No files bigger than {fmt_size(threshold)}"))
        return

    max_size = big[0]["size"]
    print(f"\n  {c(BOLD, f'Files bigger than {fmt_size(threshold)}')} ({len(big)} files)\n")
    for f in big[:50]:
        b = bar(f["size"], max_size, 20)
        rel = os.path.relpath(f["path"])
        print(f"  {fmt_size(f['size']):>10}  {b}  {rel}")
    if len(big) > 50:
        print(c(DIM, f"  ... and {len(big) - 50} more"))
    total = sum(f["size"] for f in big)
    print(f"\n  {c(DIM, f'Total: {fmt_size(total)}')}\n")


def cmd_smaller(files: list[dict], threshold: int):
    """Show files smaller than threshold."""
    small = [f for f in files if f["size"] < threshold]
    print(f"\n  {c(BOLD, f'Files smaller than {fmt_size(threshold)}')}: "
          f"{len(small)} files, {fmt_size(sum(f['size'] for f in small))} total\n")


def cmd_older(files: list[dict], days: int):
    """Show files not modified in N days."""
    cutoff = time.time() - (days * 86400)
    old = [f for f in files if f["mtime"] < cutoff]
    old.sort(key=lambda f: f["size"], reverse=True)

    if not old:
        print(c(DIM, f"  No files older than {days} days"))
        return

    print(f"\n  {c(BOLD, f'Files not modified in {days} days')} ({len(old)} files)\n")
    for f in old[:30]:
        age = int((time.time() - f["mtime"]) / 86400)
        rel = os.path.relpath(f["path"])
        print(f"  {fmt_size(f['size']):>10}  {age:>4}d ago  {rel}")
    if len(old) > 30:
        print(c(DIM, f"  ... and {len(old) - 30} more"))
    total = sum(f["size"] for f in old)
    print(f"\n  {c(DIM, f'Total: {fmt_size(total)}')}\n")


def output_json(files: list[dict], path: str):
    """Output as JSON."""
    total_size = sum(f["size"] for f in files)
    by_ext: dict[str, int] = defaultdict(int)
    for f in files:
        by_ext[f["ext"]] += f["size"]

    result = {
        "path": path,
        "total_bytes": total_size,
        "total_human": fmt_size(total_size),
        "file_count": len(files),
        "by_extension": dict(sorted(by_ext.items(), key=lambda x: -x[1])),
        "top_files": [
            {"path": f["path"], "size": f["size"], "ext": f["ext"]}
            for f in sorted(files, key=lambda f: f["size"], reverse=True)[:20]
        ],
    }
    print(json.dumps(result, indent=2))


def main():
    parser = argparse.ArgumentParser(
        description="diskuse -- analyze disk usage",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("path", nargs="?", default=".", help="Directory to analyze (default: .)")
    parser.add_argument("--top", type=int, metavar="N", help="Show top N largest files")
    parser.add_argument("--ext", action="store_true", help="Breakdown by file extension")
    parser.add_argument("--dirs", action="store_true", help="Breakdown by subdirectory")
    parser.add_argument("--bigger", metavar="SIZE", help="Files bigger than SIZE (e.g. 10M, 1G)")
    parser.add_argument("--smaller", metavar="SIZE", help="Files smaller than SIZE")
    parser.add_argument("--older", type=int, metavar="DAYS", help="Files not modified in DAYS")
    parser.add_argument("--format", choices=["text", "json"], default="text")
    parser.add_argument("--skip-hidden", action="store_true", help="Skip hidden files/dirs")

    args = parser.parse_args()

    target = os.path.abspath(args.path)
    if not os.path.isdir(target):
        print(f"Error: {args.path} is not a directory", file=sys.stderr)
        sys.exit(1)

    # Scan
    files = scan_directory(target, skip_hidden=args.skip_hidden)

    if args.format == "json":
        output_json(files, args.path)
        return

    # Default: show summary
    cmd_summary(files, args.path)

    # Additional views
    any_specific = False

    if args.top:
        cmd_top(files, args.top)
        any_specific = True

    if args.ext:
        n = args.top if args.top else None
        cmd_ext(files, n)
        any_specific = True

    if args.dirs:
        n = args.top if args.top else None
        cmd_dirs(files, n)
        any_specific = True

    if args.bigger:
        cmd_bigger(files, parse_size(args.bigger))
        any_specific = True

    if args.smaller:
        cmd_smaller(files, parse_size(args.smaller))
        any_specific = True

    if args.older:
        cmd_older(files, args.older)
        any_specific = True

    # Default: show dirs + top 10 + extensions
    if not any_specific:
        cmd_dirs(files, 15)
        cmd_top(files, 10)
        cmd_ext(files, 10)


if __name__ == "__main__":
    main()