swap_report.py

#!/usr/bin/env -S uv run --quiet --script
# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "rich",
# ]
#
# ///

"""
swap_report.py - See how swap is being used by what processes
Python version by: Michael Garcia <thecrazygm@gmail.com> 2025-06-04
"""

import argparse
import os
import pwd
import re
from collections import namedtuple
from datetime import datetime

from rich import box
from rich.console import Console
from rich.table import Table
from rich.text import Text

Process = namedtuple("Process", ["pid", "swap_kb", "user", "cmd"])


def human_readable(kb):
    if kb >= 1048576:
        return f"{kb / 1048576:.2f} GiB"
    elif kb >= 1024:
        return f"{kb / 1024:.2f} MiB"
    else:
        return f"{kb} kB"


def get_processes(user_filter=None, name_filter=None):
    processes = []
    for pid in filter(str.isdigit, os.listdir("/proc")):
        smaps_path = f"/proc/{pid}/smaps"
        if not os.path.exists(smaps_path):
            continue
        try:
            with open(smaps_path) as f:
                swap_kb = sum(
                    int(line.split()[1]) for line in f if line.startswith("Swap:")
                )
            if swap_kb == 0:
                continue
        except Exception:
            continue
        try:
            user = pwd.getpwuid(int(os.stat(f"/proc/{pid}").st_uid)).pw_name
        except Exception:
            user = "?"
        if user_filter and user != user_filter:
            continue
        try:
            with open(f"/proc/{pid}/cmdline", "rb") as f:
                cmd = f.read().replace(b"\0", b" ").decode().strip()
            if not cmd:
                with open(f"/proc/{pid}/comm") as f2:
                    cmd = f2.read().strip()
        except Exception:
            cmd = "?"
        if name_filter and not re.search(name_filter, cmd):
            continue
        processes.append(Process(pid, swap_kb, user, cmd))
    return processes


def color_for_swap(kb):
    if kb >= 1048576:
        return "bold white on red"
    elif kb >= 102400:
        return "bold black on yellow"
    else:
        return ""


def main():
    parser = argparse.ArgumentParser(description="Show swap usage by process.")
    parser.add_argument(
        "-n", "--num", type=int, default=30, help="Show top N processes (default: 30)"
    )
    parser.add_argument(
        "-s",
        "--sort",
        choices=["swap", "pid", "name"],
        default="swap",
        help="Sort by (swap, pid, name)",
    )
    parser.add_argument("-u", "--user", help="Filter by user")
    parser.add_argument("-p", "--pattern", help="Filter by process name (regex)")
    parser.add_argument("-o", "--output", help="Output to file")
    parser.add_argument("--no-color", action="store_true", help="Disable color output")
    args = parser.parse_args()

    procs = get_processes(user_filter=args.user, name_filter=args.pattern)
    if args.sort == "swap":
        procs.sort(key=lambda p: p.swap_kb, reverse=True)
    elif args.sort == "pid":
        procs.sort(key=lambda p: int(p.pid))
    elif args.sort == "name":
        procs.sort(key=lambda p: p.cmd.lower())

    console = Console(record=bool(args.output))
    table = Table(
        title=f"Swap Usage Report - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        box=box.SIMPLE_HEAVY,
    )
    table.add_column("PID", justify="right", style="cyan", no_wrap=True)
    table.add_column("Swap", justify="right", style="magenta")
    table.add_column("User", style="green")
    table.add_column("Process", style="white")
    total_swap = 0
    for proc in procs[: args.num]:
        swap_str = human_readable(proc.swap_kb)
        row_style = color_for_swap(proc.swap_kb) if not args.no_color else ""
        table.add_row(str(proc.pid), swap_str, proc.user, proc.cmd, style=row_style)
        total_swap += proc.swap_kb
    console.print(table)
    summary = Text(
        f"Total swap used by listed processes: {human_readable(total_swap)}",
        style="bold",
    )
    console.print(summary)
    if args.output:
        console.save_html(args.output)


if __name__ == "__main__":
    main()