mlk

from peak.util.assembler import *
from peak.util.decorators import rewrap
import new
import inspect


def accepts(*argtypes):
    def decorate(func):
        argnames = inspect.getargspec(func)[0]
        code_args = [process_arg(i, argnames[i], cls) for i, cls in enumerate(argtypes)]
        c = Code.from_function(func)
        c.return_(Call(Const(func), code_args))
        wrapped = new.function(c.code(), {})
        return rewrap(func, wrapped)
    return decorate


@nodetype()
def process_arg(i, argname, spec, code=None):
    if code is None:
        return i, argname, spec
    else:
        arg = Local(argname)
        code(Call(Const(isinstance), [arg, Const(spec)]))
        skip = code.JUMP_IF_TRUE()
        error_msg = Suite([ Const("Expected %r for argument #%d got %%r" % (spec, i+1)),
                            Call(Const(type), [arg]),
                            Code.BINARY_MODULO])
        code(Call(Const(TypeError), [error_msg]))
        code.RAISE_VARARGS(1)
        skip()
        code.POP_TOP()
        code(arg)


@accepts(float, float)
def float_div(a, b):
    """
    A function that divides floats.
    Any other types of arguments will throw an exception

        >>> float_div(1.0, 2.0)
        0.5
        >>> float_div(1, 1.0)
        Traceback (most recent call last):
        ...
        TypeError: Expected <type 'float'> for argument #1 got <type 'int'>
    """
    return a / b


if __name__ == '__main__':
    import doctest
    doctest.testmod(verbose=True)