Search code examples
pythonpython-3.xpython-decorators

Double wrapping Fibonacci


I try to use two different wrappers on a recursive Fibonacci function to: 1) Count the number of recursions, 2) Memoize the computed values to reduce computations needed.

As each wrapping function creates a new function attribute on the created function, after wrapping again I can't access it anymore. Is there a way to still access it without doing a single wrap function with both effects ?

def count(f):
    def f1(*args):
        f1.counter += 1
        return f(*args)
    f1.counter = 0
    return f1

def memoize(f):
    def f1(*args):
        if not args in f1.memo:
            f1.memo[args] = f(*args)
        return f1.memo[args]
    f1.memo = {}
    return f1

@memoize
@count
def fib(n):
    if n < 2:
        return n
    return fib(n - 1) + fib(n - 2)

After this I can access fib.memo but not fib.counter, and conversely if I wrap fib with memoize before count.


Solution

  • This question was originally closed as a duplicate to How to make a chain of function decorators?, where it is correctly pointed out that using functools.wraps to chain decorators can help by copying all the attributes of the wrapped function to the wrapper function.

    However, I'm reopening this question because using functools.wraps won't completely work in your case due to your counter attribute being an immutable integer, so having wraps copy the counter attribute to the wrapper actually only copies its initial value of 0 to the wrapper, and subsequent changes to the counter attribute of the wrapped function won't be reflected in the wrapper's counter attribute, which is why:

    from functools import wraps
    
    def count(f):
        @wraps(f)
        def f1(*args):
            f1.counter += 1
            return f(*args)
        f1.counter = 0
        return f1
    
    def memoize(f):
        @wraps(f)
        def f1(*args):
            if not args in f1.memo:
                f1.memo[args] = f(*args)
            return f1.memo[args]
        f1.memo = {}
        return f1
    
    @memoize
    @count
    def fib(n):
        if n < 2:
            return n
        return fib(n - 1) + fib(n - 2)
    
    print(fib(4))
    print(fib.memo)
    print(fib.counter)
    

    incorrectly outputs:

    3
    {(1,): 1, (0,): 0, (2,): 1, (3,): 2, (4,): 3}
    0
    

    To remedy this, you have to initialize counter with a mutable object. Since Python does not have a mutable integer natively, you can create a custom class that wraps around an integer instead:

    class Int:
        def __init__(self, value=0):
            self.value = value
    
    def count(f):
        @wraps(f)
        def f1(*args):
            f1.counter.value += 1
            return f(*args)
        f1.counter = Int()
        return f1
    

    so that:

    print(fib(4))
    print(fib.memo)
    print(fib.counter.value)
    

    would correctly output:

    3
    {(1,): 1, (0,): 0, (2,): 1, (3,): 2, (4,): 3}
    5