Helps build command line tools in Python.

cli.py

# -*- coding: utf-8 -*-

"""A simple tool for making command line tools in python."""

import os
import sys


class CLI(object):

    HELP_OPTION = "--help"
    
    TEXT_BOLD = '\033[01m'
    TEXT_REVERSE = '\033[07m'
    TEXT_DISABLE = '\033[02m'
    TEXT_UNDERLINE = '\033[04m'
    TEXT_STRIKETHROUGH = '\033[09m'

    TEXT_OK = '\033[92m'
    TEXT_WARNING = '\033[93m'
    TEXT_FAIL = '\033[91m'

    TEXT_BG_OK = '\033[42m'
    TEXT_BG_WARNING = '\033[43m'
    TEXT_BG_FAIL = '\033[41m'

    TEXT_RESET = '\033[0m'

    text_fail = [TEXT_BOLD, TEXT_FAIL]

    commands = {}
    execution_stack = []

    @classmethod
    def format_for_tty(cls, text, formats):
        """Format text for output to a TTY."""
        pre = "".join(formats) if formats else ""
        post = cls.TEXT_RESET if formats else ""
        return pre + text + post

    @classmethod
    def echo(cls, text, formats=None):
        """Safely echo output to STDOUT."""
        output = text
        if sys.stdout.isatty():
            output = cls.format_for_tty(text, formats)
        sys.stdout.write(output + os.linesep)

    @classmethod
    def error(cls, text, formats=None):
        """Safely echo error to STDERR, and exit with a status code."""
        output = text
        if sys.stderr.isatty():
            output = cls.format_for_tty(text, formats)
        sys.stderr.write(output + os.linesep)

    @classmethod
    def exit(cls, code=1):
        """Exit cleanly with a correct code."""
        sys.exit(code)

    @classmethod
    def find_by_alias(cls, alias, collection):
        """Is there an item with the specified collection?"""
        result = None
        for key, item in collection.items():
            if item.get("alias") == alias:
                result = key
                break
        return result
    
    @classmethod
    def init_command(cls, name):
        """Get or create a command from the ``commands`` dictionary."""
        if name not in cls.commands:
            cls.commands[name] = {
                "handler": None,
                "alias": None,
                "arguments": [],
                "options": {},
                }
        return cls.commands[name]

    @classmethod
    def set_alias(cls, command, alias):
        """Set an alias for a command."""
        cls.commands[command]["alias"] = alias

    @classmethod
    def set_handler(cls, command, handler):
        """Set a handler for a command."""
        cls.commands[command]["handler"] = handler

    @classmethod
    def get_handler(cls, command):
        """Get a handler for a command."""
        return cls.commands[command]["handler"]

    @classmethod
    def set_argument(cls, command, argument):
        """Set an argument for a command."""
        argument_data = {
            "name": argument,
            }
        cls.commands[command]["arguments"].append(argument_data)

    @classmethod
    def get_arguments(cls, command):
        """Get all arguments for a command."""
        return cls.commands[command]["arguments"]

    @classmethod
    def set_option(cls, command, option, alias, default=None, help=None):
        """Set an option for a command."""
        option_data = {
            "name": option,
            "alias": alias,
            "default": default,
            "help": help,
            }
        cls.commands[command]["options"][option] = option_data

    @classmethod
    def get_options(cls, command):
        """Get all options for a command."""
        return cls.commands[command]["options"]

    @classmethod
    def get_options_defaults(cls, options):
        """Get a dict of options defaults."""
        result = {}
        for key, item in options.items():
            default = item.get("default")
            if default:
                result[key] = default
        return result

    @classmethod
    def get_command_parts(cls, alias):
        """Get the command handler, arguments, and options."""
        command = cls.find_by_alias(alias, cls.commands)
        if not command:
            cls.error("No command: " + str(alias), formats=cls.text_fail)
            cls.exit()
        handler = cls.get_handler(command)
        arguments = cls.get_arguments(command)
        options = cls.get_options(command)
        return (handler, arguments, options)

    @classmethod
    def command(cls, alias=None):
        """Register a function, optionally under an alias."""
        def wrapper(func):
            command = func.__name__
            cls.init_command(command)
            cls.set_alias(command, alias)
            cls.set_handler(command, func)
            return func
        return wrapper

    @classmethod
    def argument(cls, argument):
        """Register an argument for a command."""
        def wrapper(func):
            command = func.__name__
            cls.init_command(command)
            cls.set_argument(command, argument)
            return func
        return wrapper

    @classmethod
    def option(cls, option, alias=None, default=None, help=None):
        """Register an option for a command."""
        def wrapper(func):
            command = func.__name__
            cls.init_command(command)
            cls.set_option(command, option, alias, default, help)
            return func
        return wrapper

    @classmethod
    def show_help(cls, alias):
        """Display help for a command."""
        handler, arguments, options = cls.get_command_parts(alias)
        args = [str(x.get("name")).upper() for x in arguments]
        pprint_args = " ".join(args)
        cls.echo("USAGE: " + alias + " [OPTIONS] " + pprint_args)
        if handler.__doc__:
            cls.echo("")
            cls.echo("  " + handler.__doc__)
        cls.echo("")
        cls.echo("OPTIONS")
        if options:
            for option in options.values():
                option_text = ""
                alias_text = option.get("alias")
                if alias_text:
                    option_text += alias_text
                help_text = option.get("help")
                if help_text:
                    option_text += "\t" + help_text
                cls.echo(option_text)
        cls.echo(cls.HELP_OPTION + "\t" + "Display help.")
        sys.exit()

    @classmethod
    def parse_command(cls, alias, args):
        """Parse the provided args for a command."""
        handler, arguments, options = cls.get_command_parts(alias)

        # We'll populate these as we proceed.
        i = 0
        prepared_arguments = []
        prepared_options = cls.get_options_defaults(options)

        # Go through the provided args, one by one.
        while i < len(args):

            # Is the current arg an alias for an option or argument?
            arg = args[i]

            # Show help if it's the help option.
            if arg == cls.HELP_OPTION:
                cls.show_help(alias)

            # If it's an alias for an option, add the option and its value
            # to the list of prepared options.
            option = cls.find_by_alias(arg, options)
            if option:
                option_name = options[option]["name"]
                i += 1
                if len(args) > i:
                    option_value = args[i]
                else:
                    cls.error("Missing value for " + arg, formats=cls.text_fail)
                    cls.exit()
                option_value = args[i]
                prepared_options[option_name] = option_value
                i += 1

            # If it's not an alias for an option, it must be an argument. If we're
            # expecting an argument, add it to the list of prepared arguments.
            elif len(prepared_arguments) < len(arguments):
                prepared_arguments.append(arg)
                i += 1

            # Otherwise, we don't recognize this as an alias for anything.
            else:
                break

        # Are we missing arguments now?
        if len(prepared_arguments) < len(arguments):
            missing_args = arguments[len(prepared_arguments):]
            pprint_list = [x.get("name") for x in missing_args]
            msg = "Missing argument(s): " + ", ".join(pprint_list)
            cls.error(msg, formats=cls.text_fail)
            cls.exit()

        # We're good. Add the handler and its arguments/options to the execution
        # stack so it can get executed after we finish parsing the commands.
        else:
            execution_data = {
                "handler": handler,
                "arguments": prepared_arguments,
                "options": prepared_options
                }
            cls.execution_stack.append(execution_data)

        return i + 1

    @classmethod
    def build_execution_stack(cls, args):
        """Build the execution stack."""
        i = 0
        while i < len(args):
            next_arg = cls.parse_command(args[i], args[i + 1:])
            i += next_arg

    @classmethod
    def run_execution_stack(cls):
        """Run the execution stack."""
        for command in cls.execution_stack:
            handler = command.get("handler")
            arguments = command.get("arguments")
            options = command.get("options")
            handler(*arguments, **options)

    @classmethod
    def run(cls, args):
        """Take a string of arguments, parse them, and run the program."""
        cls.build_execution_stack(args)
        cls.run_execution_stack()


@CLI.command(alias="hi")
@CLI.argument("text")
@CLI.option("is_red", alias="--red", default=False, help="Display in red.")
@CLI.option("punctuation", alias="--punc", help="Punctuation for the end.")
def greeting(text, is_red, punctuation=None):
    """Print out a simple greeting."""
    message = text
    if punctuation:
        message += punctuation
    if is_red:
        formats = [CLI.TEXT_FAIL, CLI.TEXT_BOLD]
    else:
        formats = [CLI.TEXT_OK]
    CLI.echo(message, formats=formats)


if __name__ == "__main__":
    specified_arguments = sys.argv[1:]
    CLI.run(specified_arguments)