pyaz.py

import logging
import shlex
from collections import namedtuple
from io import StringIO

from azure.cli.core import get_default_cli
from knack.log import CLI_LOGGER_NAME

_cli = get_default_cli()
AzResult = namedtuple('AzResult', ['exit_code', 'out', 'log'])

# adjust the logging level if you want INFO or DEBUG log
logging_level = logging.WARNING


def az(args_str):
    """
    Invoke Azure CLI command with a Python function call. For example:
        az("group show -n mygroup")

    exit_code can be one of
        0: succeed
        1: request error
        2: command parsing error
        3: resource doesn't exist

    :param args_str: the argument string, without az
    :return: a named tuple consisting of exit_code, out and log
    """
    stdout_buf = StringIO()
    log_buf = StringIO()

    # root logger only keeps CRITICAL
    root_logger = logging.getLogger()
    root_handler = logging.StreamHandler(stream=log_buf)
    root_handler.setLevel(logging.CRITICAL)
    root_logger.addHandler(root_handler)

    # cli logger keeps WARNING and above
    cli_logger = logging.getLogger(CLI_LOGGER_NAME)
    cli_handler = logging.StreamHandler(stream=log_buf)
    cli_handler.setLevel(logging_level)
    cli_logger.addHandler(cli_handler)

    args = shlex.split(args_str)
    try:
        # exit code 0 or 1
        exit_code = _cli.invoke(args, out_file=stdout_buf)
    except SystemExit as ex:
        # exit code 2 or 3, generated by sys.exit()
        exit_code = ex.code

    root_logger.removeHandler(root_handler)
    cli_logger.removeHandler(cli_handler)

    return AzResult(exit_code, stdout_buf.getvalue(), log_buf.getvalue())