Pretty Print for Slurm `squeue`. It uses [rich](https://github.com/Textualize/rich). You can install it globally via `sudo pip install rich`.

psqueue.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""psqueue.py
This script displays squeue in a pretty table.
"""
import os
import subprocess
import sys

try:
    from rich.console import Console
    from rich.table import Table
except ImportError:
    sys.exit("Please install rich as sudo: sudo pip install rich")


try:
    command = ['sinfo', '--format=%N', '--noheader']
    result = subprocess.run(command, check=True, stdout=subprocess.PIPE, text=True)
    NODE_NAMES = [x.strip() for x in result.stdout.split(',')]
except subprocess.CalledProcessError as e:
    NODE_NAMES = []


def run_squeue(args):
    # Command construction
    command = ['squeue'] + args
    try:
        # Running squeue command
        result = subprocess.run(command, check=True, stdout=subprocess.PIPE, text=True)
        return result.stdout
    except subprocess.CalledProcessError as e:
        print(f"Error running squeue: {e}")
        sys.exit(1)


def parse_output(output):
    # Splitting the output into lines
    lines = output.strip().split('\n')
    # Assuming first line is header
    headers = lines[0].split()

    # Get the index of the "NODELIST" column
    nodelist_index = headers.index("NODELIST") if "NODELIST" in headers else None

    # Parsing each row
    rows = []
    for line in lines[1:]:
        parts = line.split()
        # in case the "NodeList" column is not present, add it
        if nodelist_index is not None:
            nodelist_str = parts[nodelist_index]
            # dirty hack to add "None" to the nodelist column
            if nodelist_str.strip() not in NODE_NAMES:
                parts.insert(nodelist_index, "None")
        rows.append(parts)
    return headers, rows


def get_state_color(state):
    if state in ["COMPLETED", "CD"]:
        return "white"
    elif state in ["RUNNING", "COMPLETING", "R", "CG"]:
        return "green"
    elif state in ["PENDING", "PD"]:
        return "cyan"
    return "bright_red"


def get_time_left_color(time_left):
    def split_time(t):
        if ':' in t:
            p = t.split(':')
            if len(p) == 3:
                return int(p[0]), int(p[1]), int(p[2])
            elif len(p) == 2:
                return 0, int(p[0]), int(p[1])
        return 0, 0, int(t)

    if '-' in time_left:
        days, time = time_left.split('-')
        days = int(days)
        hours, minutes, seconds = split_time(time)
    else:
        days = 0
        hours, minutes, seconds = split_time(time_left)

    if days >= 1:
        return "white"
    elif hours >= 12:
        return "navajo_white1"
    elif hours >= 1:
        return "light_salmon1"
    elif minutes >= 30:
        return "red1"
    return "bright_red"


def display_table(headers, rows, title="squeue"):
    console = Console()
    table = Table(title=title, show_header=True, highlight=True)

    # merge TRES_PER_JOB and TRES_PER_NODE
    if "TRES_PER_JOB" in headers and "TRES_PER_NODE" in headers:
        index_job = headers.index("TRES_PER_JOB")
        index_node = headers.index("TRES_PER_NODE")
        for row in rows:
            if row[index_job] == "N/A":
                row[index_job] = row[index_node]
            row.pop(index_node)
        headers.pop(index_node)

    # if the number of columns of a row is greater than the number of headers, concat the last columns
    for i, row in enumerate(rows):
        if len(row) > len(headers):
            rows[i] = row[:len(headers)-1] + [' '.join(row[len(headers)-1:])]

    # Adding columns to the table
    for header in headers:
        if header == "TRES_PER_JOB":
            table.add_column("GPUS", no_wrap=True)
        else:
            table.add_column(header, no_wrap=True)

    # Get the current user's name to highlight it
    current_user = os.getenv("USER")

    # Get the index of the "USER" column
    user_index = headers.index("USER") if "USER" in headers else None
    reason_index = headers.index("REASON") if "REASON" in headers else None
    state_index = headers.index("ST") if "ST" in headers else None
    state_index = headers.index("STATE") if "STATE" in headers else state_index
    tres_index = headers.index("TRES_PER_JOB") if "TRES_PER_JOB" in headers else None

    # timeleft_index = headers.index("TIME_LEFT") if "TIME_LEFT" in headers else None

    # Adding rows to the table
    for row in rows:
        # Replace "gres:gpu:" with ""
        if tres_index is not None:
            row[tres_index] = row[tres_index].replace("gres:gpu:", "")

        # Replace "None" with empty string
        if reason_index is not None and row[reason_index].lower() == "none":
            row[reason_index] = ""

        # Highlight the state
        if state_index is not None:
            state = row[state_index]
            color = get_state_color(state)
            row[state_index] = f"[{color}]{state}[/]"

        # Highlight the time left
        # if timeleft_index is not None:
        #     time_left = row[timeleft_index]
        #     color = get_time_left_color(time_left)
        #     row[timeleft_index] = f"[{color}]{time_left}[/]"

        # Highlight the current user's name
        style = None
        if user_index is not None and row[user_index] == current_user:
            style = "bold"

        table.add_row(*row, style=style)

    # Displaying the table
    console.print(table)


if __name__ == "__main__":
    # Getting additional arguments passed to the script
    additional_args = sys.argv[1:]

    # If no specific format is provided, use the default format
    args_str = ' '.join(additional_args)
    if '--Format' not in args_str and '-O' not in args_str:
        additional_args += ["--Format=JobID,Partition,Name,UserName,QOS,TimeUsed,TimeLeft,NumCPUs,tres-per-job,tres-per-node,MinMemory,NodeList,State,Reason:64"]  # noqa

    # if '--format' not in args_str and '-o' not in args_str:
    #     additional_args += ["--format=%.7i %9P %35j %.8u %.12q %.12N %.12M %.12L %.5C %.7m %.4D %.12T %r"]

    output = run_squeue(additional_args)
    headers, rows = parse_output(output)
    display_table(headers, rows, title="squeue "+args_str)

Example: