cli.py

"""
Inspect functions or objects containing funcs to produce CLIs with argparse.

The meat of this module is generate_parser and (to a lesser extent)
generate_parser_obj.

Currently there is no easy way to add documentation to parameters.

"""

import argparse
import inspect
import sys

def cli(obj, *, default_type=None):
    """
    Rtn a function that, when called with sys.argv will act like a cli to obj.

    The wrapper performs simple validation and conversion on argv, then calls
    func. For details of signature interpretation, see generate_parser.

    """
    if inspect.isfunction(obj):
        return function2cli(obj, default_type=default_type)
    else:
        return obj2cli(obj, default_type=default_type)

def function2cli(func, *, default_type=None):
    """
    Rtn a function that, when called with sys.argv will act like a cli to func.

    The wrapper performs simple validation and conversion on argv, then calls
    func. For details of signature interpretation, see generate_parser.

    """
    parser = generate_parser(func, default_type=default_type)

    def inner(argv=None, exit=True):
        import sys
        if argv is None:
            argv = sys.argv[1:]

        result = apply_namespace(func, parser.parse_args(argv))

        if result is not None:
            print(result)
        if exit: sys.exit(0)

    return inner


def obj2cli(obj, *, default_type=None):
    """
    Rtn a cli to a module, class, or other entity having function attributes.

    Each of the function attributes of obj are added as subparsers.

    See also: https://pypi.org/project/cli2/

    """
    parser = generate_parser_obj(obj, default_type=default_type)

    def inner(argv=None, exit=True):
        import sys
        if argv is None:
            argv = sys.argv[1:]

        args = parser.parse_args(argv)
        # Delete subparser identifier
        args.__delattr__('{command}')
        func = args.__getattribute__('{func}')
        args.__delattr__('{func}')

        result = apply_namespace(func, args)

        if result is not None:
            print(result)
        if exit: sys.exit(0)

    return inner


def opportunistic(coercion):
    "Make a coercion function opportunistic - converts if it can, leaves as is otherwise"
    def inner(value):
        try:
            return coercion(value)
        except:
            return value

    return inner


def coerce_bool(string):
    if string.lower() in ('y', 'yes', 'true', '1'):
        return True
    elif string.lower() in ('n', 'no', 'false', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError(f'Cannot convert "{string}" to bool.')


def coerce_number(value):
    try: return int(value) 
    except ValueError:
        try: return float(value)
        except ValueError:
            try: return complex(value)
            except ValueError:
                raise argparse.ArgumentTypeError(f'Cannot coerce "{value}" to int, float or complex')


def _coerce_numbers(bind):
    "Coerce the values of unannotated params to numbers, if possible."

    def coerce_one(value):
        try: value = int(value) 
        except ValueError:
            try: value = float(value)
            except ValueError:
                try: value = complex(value)
                except ValueError:
                    pass
        return value

    for value, param in zip(bind.arguments.values(), bind.signature.parameters.values()):
        if _isempty(param.annotation) and value is not None:
            bind.arguments[param.name] = coerce_one(value)

    # bind is edited in-place. Return it as well for ease of use.
    return bind


def generate_parser_obj(obj, *, default_type=None):
    """
    argparse parser automatically generated by inspecting an obj and its functions.
    
    """
    parser = argparse.ArgumentParser(
            prog=obj.__name__,
            # Modules tend to have very long docstrings...
            description=inspect.cleandoc(obj.__doc__).splitlines()[0] if obj.__doc__ else None
            )
    # Dest is suppressed by default. Set it to a value that no parameter name can take.
    subparsers = parser.add_subparsers(dest='{command}')
    # Due to a stupid bug, subparsers are optional by default.
    subparsers.required = True

    def functions(obj):
        return tuple([o[1] for o in inspect.getmembers(obj) if inspect.isfunction(o[1])])

    for func in functions(obj):
        thisparser = subparsers.add_parser(
                func.__name__,
                description=inspect.cleandoc(func.__doc__) if func.__doc__ else None
                )
        generate_parser(func, thisparser, default_type=default_type)
        # Save a reference for obj2cli to call if this subparser is used.
        thisparser.set_defaults(**{'{func}':func})

    return parser

def generate_parser(func, parser=None, *, default_type=None):
    """
    argparse parser automatically generated by inspecting func.

    Function signature interpretation:
        - (`POSITIONAL_ONLY`, `POSITIONAL_OR_KEYWORD`) = positional
        - positional with default = optional positional
        - `KEYWORD_ONLY` = options
        - defaults = defaults
            - Boolean special casing
                - If the default is `True` or `False`, the option does not take any
                  arguments. Instead, if the option is given on the
                  commandline, the opposite value to the default is given to
                  the function.
                - Example:
                    ```python
                    def rm(*, force=False):
                        pass

                    cli(rm)(['--force']) # ~== rm(force=True)
                    cli(rm)([]) # ~== rm(force=False)
                    ```
        - type annotations = type
            - If the `type` is callable, it is called by `argparse` on the relevant substring
            - If the type is `bool`, it is replaced by `coerce_bool`
            - Provide your own custom function or handle the strings in your
              function body if you need something fancier.

    Special types:
        - cli.Choice
            - `def foo(x:Choice(1,2))` interpreted as `add_argument('x', choices=(1,2), type=int)`
    """
    sig = inspect.signature(func)

    if parser is None:
        parser = argparse.ArgumentParser(
                prog=func.__name__,
                description=inspect.cleandoc(func.__doc__) if func.__doc__ else None
                )

    for param in sig.parameters.values():
        kwargs = {}

        # Positional, keyword or vararg?
        if param.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD):
            name = param.name
            if not _isempty(param.default):
                kwargs["nargs"] = '?'
        elif param.kind is inspect.Parameter.KEYWORD_ONLY:
            if len(param.name) == 1:
                name = f"-{param.name}"
            else:
                name = f"--{param.name}"
            if _isempty(param.default):
                kwargs["required"] = True
        elif param.kind is inspect.Parameter.VAR_POSITIONAL:
            name = param.name
            kwargs["nargs"] = "*"
        else:
            raise NotImplementedError(f"Params of {param.kind} are not supported yet.")

        # Argument type and default
        if not _isempty(param.default):
            kwargs["default"] = param.default
            if param.default is not None:
                kwargs["type"] = type(param.default)
        if not _isempty(param.annotation):
            if type(param.annotation) is type and issubclass(param.annotation, Choice):
                kwargs["choices"] = param.annotation.choices
                kwargs["type"] = param.annotation.type
            else:
                kwargs["type"] = param.annotation
        if "type" not in kwargs and default_type is not None:
            kwargs["type"] = default_type

        # Special casing for booleans
        if "type" in kwargs and kwargs["type"] is bool:
            if kwargs.get("default", None) in (True, False):
                # Interpret this as a flag that is true iff present
                kwargs.pop('nargs', None)
                kwargs.pop('type', None)
                kwargs["action"] = 'store_false' if kwargs["default"] is True else 'store_true'
            else:
                # Use a custom converter
                kwargs["type"] = coerce_bool

        parser.add_argument(name, **kwargs)

    return parser


def namespace_to_bind(namespace, sig):
    "argparse.Namespace -> inspect.BoundArguments"
    bind = sig.bind_partial()
    for paramname in sig.parameters:
        bind.arguments[paramname] = namespace.__getattribute__(paramname)
    return bind


def apply_namespace(func, namespace):
    "Call (or apply) func with the arguments in namespace"
    bind = namespace_to_bind(namespace, inspect.signature(func))
    return func(*bind.args, **bind.kwargs)


def _isempty(thing):
    return thing is inspect.Parameter.empty


class Choice:
    """
    Use this type to annotate parameters that should accept only some values.

    e.g.
        def dispatch(cmd:Choice('spam','shutdown','banana')='spam'):
            if cmd == 'spam:
                ...

    The argparse `type` kwarg will be set to the type of the first argument. It
    is an error to provide choices with different types.

    "Note that inclusion in the choices container is checked after any type
    conversions have been performed"
    https://docs.python.org/3/library/argparse.html#choices

    """
    def __new__(self, *choices):
        _type = type(choices[0])
        for choice in choices:
            assert type(choice) == _type, "All choices must share the same type."
        return type(f'Choice{choices}', (Choice,), {'choices':choices, 'type':_type})