from decorator import decorator
from functools import partial
from inspect import getargspec


def better_update_wrapper(wrapper,template):
    """Like functools.update_wrapper, but changes signature as well.
    Like decorator.decorator, but doesn't try any funny magic
    """
    def decorator_module_sucks(f,*a,**kw):
        return wrapper(*a,**kw)
    return decorator(decorator_module_sucks,template)

def dec_w_args(dec):
    """Turns a working decorator function with default keyword args into something that you can either
    apply directly as in `@mydec` or change the default args as in `@mydec(myarg="special")`. Also fixes
    function signatures.
    
    >>> @dec_w_args
    ... def tracing(fun=None, before="enter", after="leave"):
    ...     def traced(*args, **kw):
    ...         print before, fun.__name__, args, kw
    ...         fun(*args, **kw)
    ...         print after
    ...     return traced #decorator(traced, fun)
        
    >>> @tracing
    ... def t(x):
    ...     print x * x
        
    >>> t(2)
    enter t (2,) {}
    4
    leave
    
    >>> @tracing(before = "abandon hope ye who enter")
    ... def tdoom(x):
    ...     print x * x
    
    >>> tdoom(3)
    abandon hope ye who enter tdoom (3,) {}
    9
    leave
    """
    spec = getargspec(dec)
    assert len(spec.args) == len(spec.defaults) and spec.defaults[0] == None, "dec_w_args decorates decorators that look like like my_decorator(fun=None,...)"
    def optargs_decorator(f=None, *args, **kw):
        def curried_dec(inner_f):
            return better_update_wrapper(partial(dec, **kw)(inner_f),inner_f)
        if f:
            return curried_dec(f)
        return curried_dec #partial(decorator,curried_dec)
    return better_update_wrapper(optargs_decorator,dec)