#!/usr/bin/python

from __future__ import print_function

import functools
import inspect

import wrapt


def decorator_factory(*kw_only_parameters):
    """Transforms a function with keyword arguments into one with
    keyword-only arguments.

    Call this decorator as @decorator_factory() for the default mode,
    which makes all arguments with defaults keyword-only, or with the
    names of arguments to make keyword-only.  They must correspond
    with the names of arguments in the decorated function.  It works
    by collecting all the arguments into *args and **kws, then moving
    the arguments marked as keyword-only from **kws into *args.

    Args:
      *kw_only_parameters: Keyword-only arguments as strings.

    Returns:
      A decorator that modifies a function so it has keyword-only
      arguments.

    """
    def decorator(wrapped):
        """The decorator itself, assigns arguments as keyword-only and
        calculates sets for error checking.

        Args:
          wrapped: The function to decorate.

        Returns:
          A function wrapped so that it has keyword-only arguments.

        """

        # Each Python 3 argument has two independent properties: it is
        # positional-and-keyword *or* keyword-only, and it has a
        # default value or it doesn't.
        names, varargs, _, defaults = inspect.getargspec(wrapped)

        # If there are no default values getargpsec() returns None
        # rather than an empty iterable for some reason.
        if defaults is None:
            defaults = ()
        names_with_defaults = frozenset(names[len(names) - len(defaults):])
        names_to_defaults = dict(zip(reversed(names), reversed(defaults)))
        if kw_only_parameters:
            kw_only_names = frozenset(kw_only_parameters)
        else:
            kw_only_names = names_with_defaults.copy()

        @wrapt.decorator
        def wrapper(wrapped, instance, args, kws):
            """Wrapper function, checks arguments with set operations, moves args
            from **kws into *args, and then calls wrapped().

            Args:
              *args, **kws: The arguments passed to the original function.

            Returns:
              The original function's result when it's called with the
              modified arguments.

            Raises:
              TypeError: When there is a mismatch between the supplied
                and expected arguments.

            """

            new_args = []
            args_index = 0
            for name in names:
                if name in kws:
                    # Check first if there's a bound keyword for this name
                    new_args.append(kws.pop(name))
                elif name in kw_only_names:
                    # If this name is keyword-only, check for a
                    # default or raise.
                    if name in names_to_defaults:
                        new_args.append(names_to_defaults[name])
                    else:
                        _wrong_args(wrapped, names, 
                                    kw_only_names -
                                    (names_with_defaults | frozenset(kws)),
                                    'keyword-only')
                elif args_index < len(args):
                    # Check for a positional arg.
                    new_args.append(args[args_index])
                    args_index += 1
                elif name in names_to_defaults:
                    # Check for a default value.
                    new_args.append(names_to_defaults[name])
                else:
                    # No positional arg or default for this name so raise.
                    _wrong_args(wrapped, names,
                                frozenset(names) -
                                (names_with_defaults | frozenset(kws)),
                                'positional', len(args))

            if args_index != len(args) and not varargs:
                # Too many positional arguments and no varargs, so
                # raise after subtracting off the number of kw-only
                # arguments from those expected.
                raise TypeError(
                    '%s() takes %d positional arguments but %d were given' %
                    (wrapped.__name__,
                     len(names) - len(names_with_defaults | frozenset(kws)),
                     len(args)))
            else:
                # Pass the rest of the positional args, if any.
                new_args.extend(args[args_index:])

            return wrapped(*new_args, **kws)
        return wrapper(wrapped)

    return decorator


def _wrong_args(wrapped, names, missing_args, arg_type, number_of_args=0):
    """ Raise Python 3-style TypeErrors for missing arguments."""
    ordered_args = [a for a in names if a in missing_args]
    ordered_args = ordered_args[number_of_args:]
    error_message = ['%s() missing %d required %s argument' % 
                     (wrapped.__name__, len(ordered_args), arg_type)]
    if len(ordered_args) == 1:
        error_message.append(": '%s'" % ordered_args[0])
    else:
        error_message.extend(['s: ', 
                              ' '.join("'%s'" % a for a in ordered_args[:-1]),
                              " and '%s'" % ordered_args[-1]])
    raise TypeError(''.join(error_message))


if __name__ == '__main__':
    def test(f, *args, **kws):
        print(args, kws, '-> ', end='')
        try:
            f(*args, **kws)
        except TypeError as e:
            print(e.args[0])

    @decorator_factory()
    def f(a, b, c, d, e='e'):
        print(a, b, c, d, e)

    test(f)
    test(f, 0, 1, 2, 3)
    test(f, 0, 1, 2, 3, 4)

    @decorator_factory('c')
    def f(a, b, c='c', d='d', *args, **kws):
        print(a, b, c, d, args, kws)

    test(f)
    test(f, 0, 1)
    test(f, -1, b='b')
    test(f, b='b')
    test(f, 0)
    test(f, 0, 1, 2, 3, 4, 5, c='foo', d='bar', e='baz')

    @decorator_factory('b', 'c')
    def f(a, b, c='c', d='d', *args, **kws):
        print(a, b, c, d, args, kws)

    test(f, 0)
    test(f, 0, b='b')
    test(f, 0, 1, 2)
    test(f, 0, 1, 2, b='b')
    test(f, 0, 1, 2, 3, 4, 5, c='foo', d='bar', e='baz')
    test(f, 0, 1, 2, 3, 4, 5, b='foo', c='bar', e='baz')

    class C(object):
        @decorator_factory('b', 'c')
        def __init__(self, a, b, c='c', d='d', *args, **kws):
            print(a, b, c, d, args, kws)

    test(C, 0, 1, 2, b='b')
    test(C, 0, 1, 2, 3, 4, 5, b='foo', c='bar', e='baz')
