clickutil.py

import logging

import click


class ColorFormatter(logging.Formatter):
    colors = {
        "error": dict(fg="red"),
        "exception": dict(fg="red"),
        "critical": dict(fg="red"),
        "debug": dict(fg="blue"),
        "warning": dict(fg="yellow"),
    }

    def format(self, record):
        if not record.exc_info:
            level = record.levelname.lower()
            msg = record.getMessage()
            if level in self.colors:
                prefix = click.style("{}: ".format(level), **self.colors[level])
                msg = "\n".join(prefix + x for x in msg.splitlines())
            return msg
        return logging.Formatter.format(self, record)


class ClickHandler(logging.Handler):
    _use_stderr = True

    def emit(self, record):
        try:
            msg = self.format(record)
            level = record.levelname.lower()
            click.echo(msg, err=self._use_stderr)
        except Exception:
            self.handleError(record)


def get_click_logger(name: str, level=logging.INFO) -> logging.Logger:
    handler = ClickHandler()
    handler.formatter = ColorFormatter()

    logger = logging.getLogger(name)
    logger.handlers = [handler]
    logger.propagate = False
    logger.setLevel(level)

    return logger


def walk_tree(f, node, depth=0, idx=0, size=0):
    xs = list(f(node, depth, idx, size) or ())
    n = len(xs)
    for i, x in enumerate(xs):
        walk_tree(f, x, depth + 1, i, n)


def print_command_tree(command: click.Command):
    def print_node(node, depth, idx, size):
        c, is_leaf = node
        is_tail = idx == (size - 1)

        if depth == 0:
            prefix = ""
            tree_item = ""
        else:
            prefix = "    " if is_leaf else "│   "
            tree_item = "└── " if is_tail else "├── "

        line = prefix * (depth - 1) + tree_item + c.name
        if c.help:
            line += f" - {c.help}"
        click.echo(line)

        sub_commands = getattr(c, "commands", {})
        return ((x, is_tail) for _, x in sorted(sub_commands.items()))

    walk_tree(print_node, (command, False))