unity-diskquota-usage

#!/usr/bin/env python3
import os, re, sys, grp, shutil, signal

USAGE_PERCENT_RED_THRESHOLD = 75


def human_readable_size(size_in_bytes: int) -> str:
    current_size = size_in_bytes
    current_unit = "bytes"
    if current_size == 0:
        return "0 bytes"
    for new_unit in ["KB", "MB", "GB", "TB", "PB"]:
        new_size = current_size / 1000
        if new_size < 1:
            break
        else:
            current_size = new_size
            current_unit = new_unit
    return f"{current_size:.2f} {current_unit}"


def printable_length(x: str) -> int:
    assert "\n" not in x and "\t" not in x, "no newlines or tabs allowed!"
    # https://stackoverflow.com/q/14693701/12035739
    ansi_be_gone = re.compile(r"(?:\x1B[@-_]|[\x80-\x9F])[0-?]*[ -/]*[@-~]")
    return len(ansi_be_gone.sub("", x))


def fmt_table(table: list[list]) -> list[str]:
    """
    I would use tabulate but I don't want nonstandard imports
    no table headers
    allows ANSI formatted elements
    """
    output_lines = []
    assert all(len(row) == len(table[0]) for row in table), "all rows must have the same length"
    column_widths = [0] * len(table[0])
    for row in table:
        for i, element in enumerate(row):
            if printable_length(str(element)) > column_widths[i]:
                column_widths[i] = printable_length(str(element))
    column_widths = [x + 1 for x in column_widths]  # add one space in between
    for row in table:
        line = ""
        for i, value in enumerate(row):
            padding_size = column_widths[i] - printable_length(value)
            line += value + " " * padding_size
        output_lines.append(line)
    return output_lines


def red(x) -> str:
    return f"\033[0;31m{x}\033[0m"


usage = []


def print_usage_and_exit():
    if usage != []:
        for line in fmt_table(usage):
            print(line)
    sys.exit(1)


# if timed out, print whatever usage has been collected so far
# this was removed since bad NFS requires sigkill
# signal.signal(signal.SIGTERM, lambda foo, bar: print_usage_and_exit())

dirs_to_check = [os.path.expanduser("~")]  # home directory
for gid in os.getgroups():
    gr_name = grp.getgrgid(gid).gr_name
    if not gr_name.startswith("pi_"):
        continue
    for prefix in "/project", "/work":
        dir_path = os.path.join(prefix, gr_name)
        if os.path.isdir(dir_path):
            dirs_to_check.append(dir_path)

for dir_path in dirs_to_check:
    total, used, _ = shutil.disk_usage(dir_path)
    pcent_used = (used / total) * 100
    if pcent_used >= USAGE_PERCENT_RED_THRESHOLD:
        usage.append(
            [
                dir_path,
                red(human_readable_size(used)),
                red("/"),
                red(human_readable_size(total)),
                red("="),
                red(f"{(pcent_used):.0f}%"),
            ]
        )
    else:
        usage.append(
            [
                dir_path,
                human_readable_size(used),
                "/",
                human_readable_size(total),
                "=",
                f"{(pcent_used):.0f}%",
            ]
        )
print_usage_and_exit()