Search code examples
pythondoctestpython-decoratorspython-unittest

How to use doctest with a decorated function in python?


I am using a decorator:

class Memoized(object):

    __cache = {}

    def __init__(self, func):
        self.func = func
        key = (func.__module__, func.__name__)
        # print key
        if key not in self.__cache:
            self.__cache[key] = {}
        self.mycache = self.__cache[key]

    def __call__(self, *args):
        try:
            return self.mycache[args]
        except KeyError:
            value = self.func(*args)
            self.mycache[args] = value
            return value
        except TypeError:
            return self.func(*args)

    def __get__(self, obj, objtype):
       return functools.partial(self.__call__, obj)

    def reset(self):
        for v in self.__cache.itervalues():
            v.clear()

and a function:

@Memoized
def is_tile_inside_border(x, y, z, border):
    '''Checks if a tile is inside border or not
    >>> is_tile_inside_border(85,53,7,'iran')
    3
    >>> is_tile_inside_border(85,15,7,'iran')
    0
    '''
    binary_data = get_border_binary(border, z)
    return isInside((x, y), binary_data)

But when using doctest module (python mycode.py -v):

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

Python can't find the test in doc-string. I know that is a problem with decorator. But how can I fix it?

PS: functools.update_wrapper(self, func) is not working!


Solution

  • How about we side-step the problem by using a memo decorator which returns a function instead of a class instance:

    import functools
    
    def decorator(d):
        """Make function d a decorator: d wraps a function fn.
        Authors: Peter Norvig and Darius Bacon"""
        def _d(fn):
            return functools.update_wrapper(d(fn), fn)
        functools.update_wrapper(_d, d)
        return _d
    
    @decorator
    def memo(f):
        # by Peter Norvig
        """Decorator that caches the return value for each call to f(args).
        Then when called again with same args, we can just look it up."""
        cache = {}
    
        def _f(*args):
            try:
                return cache[args]
            except KeyError:
                cache[args] = result = f(*args)
                return result
            except TypeError:
                # some element of args can't be a dict key
                return f(*args)
        _f.cache = cache
        return _f
    
    @memo
    def is_tile_inside_border(x, y, z, border):
        '''Checks if a tile is inside border or not
        >>> is_tile_inside_border(85,53,7,'iran')
        3
        >>> is_tile_inside_border(85,15,7,'iran')
        0
        '''
        # binary_data = get_border_binary(border, z)
        # return isInside((x, y), binary_data)
        return True
    
    class Foo(object):
        @memo
        def bar(self, x):
            """
            >>> Foo().bar(1)
            2
            """
            return x
    
    if __name__ == '__main__':
        import doctest
        doctest.testmod()