Colorfull argparse formatter

help.png

prettyArgparse.py

import argparse
import os
import re

# 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"Here is the {bcolors.CYAN}description{bcolors.ENDC} of this script.\n" + \
                f"With this class formatter, you can do the following:\n" + \
                f"  - using newlines with '\\n' in your argparse description {bcolors.UNDERLINE}and{bcolors.ENDC} in your parameters' help messages\n" + \
                f"  - printing the default value of your {bcolors.YELLOW}parameters{bcolors.ENDC}\n" + \
                f"  - adding some colors as coded in the method `{bcolors.BOLD}format_help{bcolors.ENDC}`",
                                     formatter_class=CustomFormatter)
    parser.add_argument(
        "-a", "--arg_one", type=str, required=True, help="This is a required parameter, not like --arg_two, which is optional"
    )
    parser.add_argument(
        "-b", "--arg_two", type=str, default="valhalla", help="This is an optinal parameter.\nNot setting it will result in using the default value"
    )
    parser.add_argument(
        "-c", "--arg_three", action="store_true", help="Another type of parameter"
    )
    return parser

def sum_up(args: argparse.Namespace) -> None:
    """
    Print the sum-up of the arguments parsed by this script

    :Args:
        args: the parsed arguments.
    """
    # Params will be colorized, so size + 5 (yellow) + 4 (endc)
    size_param = max(len(param) for param in args.__dict__.keys()) + 9
    print(f'Sum-up of the script\'s {bcolors.YELLOW}arguments{bcolors.ENDC}:')
    for param, value in args.__dict__.items():
        print(f' - %{size_param}s: %s'%(f"{bcolors.YELLOW}{param}{bcolors.ENDC}", value))

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

def main() -> None:
    """
    Main function
    """
    # Parse the arguments
    args = get_args().parse_args()
    # Sum-up the arguments
    sum_up(args)

if __name__ == '__main__':
    main()