Exact same thing as running 'watch -t -c -n 1 nvidia-smi' but with some colors

pretty_nvidia_smi.py

import argparse
import os
import re
import subprocess
import time


# User-friendly console display colors
# In case of f-formating alignment, str.rjust, str.ljust or str.center,
#    be aware of their len even if you won't see then while printing
os.system("")
class bcolors:
    PINK = '\033[95m'     # len -> 5
    PURPLE = '\033[94m'   # len -> 5
    CYAN = '\033[96m'     # len -> 5
    GREY = '\033[92m'     # len -> 5
    GREEN = '\033[32m'    # len -> 5
    YELLOW = '\033[93m'   # len -> 5
    RED = '\033[91m'      # len -> 5
    ENDC = '\033[0m'      # len -> 4
    BOLD = '\033[1m'      # len -> 4
    UNDERLINE = '\033[4m' # len -> 4


# ######################################### #
#                                           #
#                   ARGS                    #
#                                           #
# ######################################### #


class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawTextHelpFormatter):
    """
    This class allows the usage of multiple formatter at once as a formatter_class of the ArgumentParser.
    It also reimplement the `format_help` method.
    """
    def format_help(self) -> str:
        """
        Format the message
        :return: the formatted message
        """
        help = self._root_section.format_help()
        if help:
            help = self._long_break_matcher.sub('\n\n', help)
            help = help.strip('\n') + '\n'
        help = re.sub(r"[Pp]ositional arguments:", f"{bcolors.RED}Positional Arguments{bcolors.ENDC}:", help)
        help = re.sub(r"[Oo]ptional arguments:", f"{bcolors.RED}Optional Arguments{bcolors.ENDC}:", help)
        # Add color around parameters
        # Regex explained:
        #  \1 -> ' ' or '['
        #  \2 -> '-' or '--'
        #  \3 -> contains only english letters or '_'
        #  \4 -> ' ' or ']' or ',' or '.'
        help = re.sub(r"([\s\[])(-{1,2})([a-zA-Z_]+)([\s\],\.])", rf"\1{bcolors.YELLOW}\2\3{bcolors.ENDC}\4", help)
        # Add color around default
        # Regex explained:
        #  \1 -> '(default: '
        #  \2 -> anything
        #  ends with ')'
        help = re.sub(r"(\(default:\s)(.*)\)", rf"{bcolors.YELLOW}\1\2){bcolors.ENDC}", help)
        return help


def get_args() -> argparse.ArgumentParser:
    """
    Get the command line arguments.
    :return: parser, the resulting argument parser.
    """
    parser = argparse.ArgumentParser(description=f"Print pretty nvidia-smi",
                                     formatter_class=CustomFormatter)
    parser.add_argument(
        "-t", "--time", type=float, default=1, help="Time between two refreshs of the output."
    )
    return parser


# ######################################### #
#                                           #
#                   MAIN                    #
#                                           #
# ######################################### #


def main(sleep_time: float):
  process = subprocess.Popen(['nvidia-smi'],
                       stdout=subprocess.PIPE, 
                       stderr=subprocess.PIPE)
  stdout, stderr = process.communicate()   
  stdout = stdout.decode("utf-8")

  stdout = re.sub(r"(\d+)MiB(\s+)/(\s+)(\d+)MiB",
                rf"{bcolors.YELLOW}\1{bcolors.ENDC}MiB\2/\3{bcolors.YELLOW}\4{bcolors.ENDC}MiB",
                stdout)

  stdout = re.sub(r"(\s+\d+\s+)([\w\/]+\s+)([\w\/]+\s+)(\d+\s+)([\w\+]+\s+)(.+)(\s\d+)MiB",
                  rf"{bcolors.BOLD}{bcolors.PURPLE}\1{bcolors.ENDC}" + \
                  rf"\2\3" + \
                  rf"{bcolors.CYAN}\4{bcolors.ENDC}" + \
                  rf"\5" + \
                  rf"{bcolors.CYAN}\6{bcolors.ENDC}" + \
                  rf"{bcolors.YELLOW}\7{bcolors.ENDC}MiB",
                  stdout)

  stdout = re.sub(r"(PID|Process name)",
                  rf"{bcolors.CYAN}\1{bcolors.ENDC}",
                  stdout)

  stdout = re.sub(r"(GPU\s+)Name",
                  rf"\1{bcolors.PURPLE}Name{bcolors.ENDC}",
                  stdout)

  stdout = re.sub(r"(GPU\s+)([^M])",
                  rf"{bcolors.BOLD}{bcolors.PURPLE}\1{bcolors.ENDC}\2",
                  stdout)

  stdout = re.sub(r"(Memory-Usage|GPU Memory|\sUsage\s)",
                  rf"{bcolors.YELLOW}\1{bcolors.ENDC}",
                  stdout)

  stdout = re.sub(r"\|(\s+)(\d+)(\s+)([\s\w]+\.{0,3})(\s+)(Off|On)(\s+)\|",
                  rf"|\1{bcolors.BOLD}{bcolors.PURPLE}\2{bcolors.ENDC}" + \
                  rf"\3{bcolors.PURPLE}\4{bcolors.ENDC}\5\6\7|",
                  stdout)

  time.sleep(sleep_time)
  os.system('clear')
  print(stdout)


if __name__ == "__main__":
  # Parse the arguments
  args = get_args().parse_args()
  try:
    while True:
      main(args.time)
  except KeyboardInterrupt:
    main(0)