Show swap usage of processes

swapstats.py

#!/usr/bin/env python3
"""
Show the swap usage of all processes in the system.
Note: This may not be reliable because of "shared pages". Source: https://www.cyberciti.biz/faq/linux-which-process-is-using-swap/
"""
import os
import sys
import humanize
import psutil
import tabulate

def get_unit_weight(b: str):
    b = b.lower()
    if b == "kb":
        return 1024
    elif b == "mb":
        return 1024 * 1024
    elif b == "gb":
        return 1024 * 1024 * 1024
    raise Exception(f"Unknown unit: {b}")

def get_process(name_or_pid):
    for proc in psutil.process_iter(attrs=['pid', 'name']):
        if name_or_pid == proc.info['name']:
            return proc
        elif name_or_pid == str(proc.info['pid']):
            return proc
    raise Exception(f"Process {name_or_pid} not found")

def get_swap(pid: int):
    with open(f"/proc/{pid}/status", "r") as file:
        for line in file:
            if line.startswith("VmSwap:"):
                s = line.split()
                value = int(s[1])
                unit = s[2]
                unit = get_unit_weight(unit)
                return value * unit
    raise Exception(f"Swap information not found for process {pid}")

def print_all_process_swap_usag():
    swap_values: list[tuple[psutil.Process, int]] = []
    for proc in psutil.process_iter(attrs=['pid', 'name']):
        try:
            swap = get_swap(proc.info['pid'])
            swap_values.append((proc, swap))  # Store PID along with swap value
        except Exception as e:
            print(f"Error: {e}", file=sys.stderr)

    swap_values.sort(key=lambda x: x[1])  # Sort by swap value

    table = []
    HEADER = ["PID", "Name", "Swap", "percent (VMS)"]
    for proc, swap in swap_values:
        table.append([
            proc.info['pid'],
            proc.info['name'],
            humanize.naturalsize(swap),
            str(round(swap / proc.memory_info().vms * 100, 2))+"%"
        ])
    table.append(HEADER) # footer

    print(tabulate.tabulate(table, headers=HEADER))

    total_swap = sum(swap for _, swap in swap_values)
    # These two values should be the same or close
    print(f"Total Swap by summing previous results: {humanize.naturalsize(total_swap)}")
    print(f"Total Swap according to system: {humanize.naturalsize(psutil.swap_memory().used)}")

if len(sys.argv) == 1:
    print_all_process_swap_usag()
elif len(sys.argv) == 2:
    proc = get_process(sys.argv[1])
    swap = get_swap(proc.info['pid'])
    print(f"Swap used by process {proc.info['name']} ({proc.info['pid']}): {humanize.naturalsize(swap)}")
else:
    print("Usage: swap.py [PID/name]")
    sys.exit(1)