Search code examples
pythonpython-decorators

Using Python decorators to take modulus of results


I'm learning decorators in Python, and I came across some trouble with a decorator I'm using:

@curry
def modulo(mod,f):
    if mod:
        @fun.wraps(f)
        def wrapper(*args, **kw):
            result = f(*args, **kw)
            return tuple(r % mod for r in result)
    else:
        return f
    return wrapper

The curry decorator is a simple currying, so I can call mod as an argument in the following:

def _fib(p=1,q=-1,mod=None):
    ''' Fibonacci sequence '''
    @modulo(mod)
    def next_fib(pair):
        x,y = pair
        return y, p*x - q*y
    
    yield from ( y for x,y in iterate(next_fib,(1,0)) )

which works and looks nice and clean. However, say I wanted another [closely related] generator for Lucas sequences:

def _luc(p=1,q=-1,mod=None):
    ''' Lucas sequence '''
    @modulo(mod)
    def next_luc(pair):
        x,y = pair
        return y, p*x - q*y
    
    yield from ( y for x,y in iterate(next_luc,(p-2,2)) )

If I call them together, I get some sort of collision:

>>> F = _fib()
>>> print(take(10,F))
[0, 1, 1, 2, 3, 5, 8, 13, 21, 34]

>>> L = _luc()
>>> print(take(10,L))
" ... TypeError: cannot unpack non-iterable function object"

Called individually they work as expected, with the correct modular terms returned.

My question twofold: is this a namespace collision where they are both referring to modulo()? How would you go about doing something like this?

helper functions:

import itertools as it
import functools as fun

def curry(f):
    argc = f.__code__.co_argcount
    f_args = []
    f_kwargs = {}
    @fun.wraps(f)
    def wrapper(*args, **kwargs):
        nonlocal f_args, f_kwargs
        f_args += args
        f_kwargs.update(kwargs)
        if len(f_args)+len(f_kwargs) == argc:
            return f(*f_args, **f_kwargs)
        else:
            return wrapper          
    return wrapper

def iterate(f,x):
    ''' x, f(x), f(f(x)), f(f(f(x))), ... '''
    return it.accumulate(it.repeat(x), lambda fx, _: f(fx))

def take(n,iterable):
    return [x for x in it.islice(iterable,n)]

Solution

  • I found that extending the original modulo wrapper does the trick! With thanks to Thomas_Breydo for suggesting to change up the curried function

    def modulo(mod):
        def wrapper(f):
            @fun.wraps(f)
            def deco(*args,**kwargs):
                if mod:
                    return tuple(n%mod for n in f(*args,**kwargs) )
                else:
                    return f(*args,**kwargs)
            return deco
        return wrapper
    

    which resolves my issue.