Memoizing decorator that can retry

Question

I have some tasks that I'd like to memoize because they connect to a rather slow network and have to wait for the data. Unfortunately this network can be a little finnicky and we get occasional connection issues. Thus I'd like if my memoizer could know to retry a certain number of times.

import functools

# Python 3.3
def enum(*sequential, **named):
    enums = dict(zip(sequential, range(len(sequential))), **named)
    return type('Enum', (), enums)  

RetryTypes = enum("TRUE", "FALSE", "REMEMBER_EXCEPTION")

class Memoizer:

    def __init__(self, retry=RetryTypes.FALSE, retry_times=0, exception_types=Exception):
        self.retry = retry
        self.retry_times = retry_times
        self.exception_types = exception_types

    def __call__(self, function):
        d = {}

        # Retry but don't store exceptions
        if self.retry is RetryTypes.TRUE and self.retry_times > 0:
            @functools.wraps(function)
            def wrapper(*args, forget=False):
                if args not in d or forget:
                    # Try n-1 times and suppress exceptions
                    for i in range(self.retry_times-1):
                        try:
                            d[args] = function(*args)
                        except self.exception_types:
                            continue
                        else:
                            break
                    else:
                        # If we have to try n times don't suppress any exception
                        d[args] = function(*args)
                return d[args]

        # Retry and store any exceptions
        elif self.retry is RetryTypes.REMEMBER_EXCEPTION:
            @functools.wraps(function)
            def wrapper(*args, forget=False):
                if args not in d or forget:
                    # If we're retrying at all, try n-1 times and suppress exceptions
                    if self.retry_times > 1:
                        for i in range(self.retry_times-1):
                            try:
                                d[args] = function(*args)
                            except self.exception_types:
                                continue
                            else:
                                break
                        else:
                            # If we have to try n times catch the exception and store it
                            try:
                                d[args] = function(*args)
                            except Exception as e:
                                d[args] = e
                    # If we're not retrying, just catch any exception and store it
                    else:
                        try:
                            d[args] = function(*args)
                        except Exception as e:
                            d[args] = e
                if isinstance(d[args], Exception):
                    raise d[args]
                else:
                    return d[args]

        # Don't retry
        else:
            @functools.wraps(function)
            def wrapper(*args, forget=False):
                if args not in d or forget:
                    d[args] = function(*args)
                return d[args]

        return wrapper

Tests:

import math
import unittest

from memoizer import Memoizer, RetryTypes


class TestNoRetry(unittest.TestCase):

    def test_no_retry_no_error(self):
        @Memoizer()
        def f(a):
            return 1745*a

        result = f(19)
        self.assertIs(result, f(19))

    def test_no_retry_with_error(self):
        @Memoizer()
        def f(a):
            raise TypeError

        with self.assertRaises(TypeError):
            f(17)

    def test_retry_no_error(self):
        @Memoizer(retry=RetryTypes.TRUE)
        def f(a):
            return 123987/a

        result = f(245)
        self.assertIs(result, f(245))

    def test_retry_no_error_retry_times(self):
        @Memoizer(retry=RetryTypes.TRUE, retry_times=2)
        def f(a):
            return 123987/a

        result = f(245)
        self.assertIs(result, f(245))

    def test_retry_error_suppressed(self):
        global time
        time = 0
        times = 2
        @Memoizer(retry=RetryTypes.TRUE, retry_times=times)
        def f(a):
            global time
            time += 1
            if time < times:
                raise TypeError
            else:
                return  math.pow(a, a)

        result = f(13)
        self.assertIs(result, f(13))

    def test_retry_other_error_not_suppressed(self):
        global time
        time = 0
        times = 2
        @Memoizer(retry=RetryTypes.TRUE, retry_times=times, exception_types=AttributeError)
        def f(a):
            global time
            time += 1
            if time < times:
                raise OSError
            else:
                return  math.pow(a, a)

        with self.assertRaises(OSError):
            f(13)

    def test_retry_too_many_errors(self):
        @Memoizer(retry=RetryTypes.TRUE, retry_times=2)
        def f(a):
            raise OSError

        with self.assertRaises(OSError):
            f(13)

    def test_no_retry_cache_errors_no_error(self):
        @Memoizer(retry=RetryTypes.REMEMBER_EXCEPTION)
        def f(a):
            return 129384716/a

        result = f(245)
        self.assertIs(result, f(245))

    def test_retry_cache_errors_no_error(self):
        @Memoizer(retry=RetryTypes.REMEMBER_EXCEPTION, retry_times=2)
        def f(a):
            return 129384716/a

        result = f(245)
        self.assertIs(result, f(245))

    def test_no_retry_cache_errors_error(self):
        @Memoizer(retry=RetryTypes.REMEMBER_EXCEPTION)
        def f(a):
            raise OSError

        error = None

        try:
            f(245)
        except OSError as e:
            error = e

        self.assertIsNot(error, None)

        try:
            f(245)
        except OSError as e:
            self.assertIs(e, error)

    def test_retry_cache_errors_error(self):
        @Memoizer(retry=RetryTypes.REMEMBER_EXCEPTION, retry_times=2)
        def f(a):
            raise OSError

        error = None

        try:
            f(245)
        except OSError as e:
            error = e

        self.assertIsNot(error, None)

        try:
            f(245)
        except OSError as e:
            self.assertIs(e, error)

    def test_retry_cache_errors_eventual_success(self):
        global time
        time = 0
        times = 2
        @Memoizer(retry=RetryTypes.REMEMBER_EXCEPTION, retry_times=2)
        def f(a):
            global time
            time += 1
            if time < times:
                raise OSError
            else:
                import random
                return a & random.getrandbits(64)

        result = f(213125)
        self.assertIs(result, f(213125))


if __name__ == '__main__':
    unittest.main()

Everything works as expected, but I'm not a huge fan of the __call__ implementation, and I feel as though I'm missing some test cases. I'm also not positive which, if any, of the test values would be cached or interned anyway. I also wonder if this might be well served by making d a weakref.WeakValueDictionary.

Answer 1

Naming

• self.retry_times
  You retry retry_times - 1 times, in all your loops.

  Consider renaming this or removing the - 1's to make it more understandable.

• forget
  This remembers, always. It just doesn't lookup.

  Consider using a temporary variable or renaming this.

• RetryTypes
  This is a constant, not a class. This should be RETRY_TYPES.

---

PEP8

• Lines are to be a maximum of 79 characters.
  The exception to this are comments at 72.

• Surround operators with one space on both sides. self.retry_times - 1.
  The exception to this is to show precedence, 2 + 2*2.

• Keep try statements as small as possible. try:a = fn(). Not try:d['a'] = fn(). This is incase d['a'] = raises an error.

• Constants use UPPER_SNAKE_CASE.
  Classes use CammelCase.
  Everything else normally uses snake_case.

Your code is really good actually.

---

Code

---

Consider changing your if self.retry is RetryTypes.TRUE and self.retry_times > 0:.

• Remove the and self.retry_times > 0:. or;

• Change it to self.retry_times > 1:.

Two small changes. The former adds readability, but makes the code worse. The latter improves the code, as if you try once then it will use the quicker wrapper.

I'm more in favour of the latter.

---

The second wrapper, is too over complicated. There is no need for the outer if else statement.

if self.retry_times > 1:
    # ...
else:
    # Duplicate code.

The else doesn't run if the for loop breaks. If the for loop never runs, it can never break. remove the if-else statement to get smaller and simpler code. Just like the first wrapper.

---

The first two wrappers are quite alike. The difference is the else statement, and the return/raise. Also if you ignore the for loop, all the wrappers are roughly the same.

You can make four functions, and select two that are needed.

# In the __call__
def plain_else(args):
    return function(*args)

def exception_else(args):
    try:
        return function(*args)
    except Exception as e:
        return e

def plain_return(r_value):
   return r_value

def exception_return(r_value):
    if isinstance(r_value, Exception):
        raise r_value
    return r_value

This allows you to change the way the function works, with no expensive if/else statement.

As my previous solution used the root of all evil, premature optimisations. This one will only use one main loop.

First we select the correct function to use. We then need to make it so if RetryTypes.FALSE is passed, it will never loop. Then we will make a simple main.

if self.retry is RetryTypes.REMEMBER_EXCEPTION:
    wrapper_else = exception_else
    wrapper_return = exception_return
else:
    wrapper_else = plain_else
    wrapper_return = plain_return

self.retry_times -= 1

if self.retry is RetryTypes.FALSE:
    self.retry_times = 0

@functools.wraps(function)
def wrapper(*args, forget=False):
    if args not in d or forget:
        for _ in range(self.retry_times):
            try:
                d[args] = function(*args)
            except self.exception_types:
                continue
            else:
                break
        else:
            d[args] = wrapper_else(args)
    return wrapper_return(d[args])

This allows the use of only one main wrapper. It handles if you want to handle errors at the end, and allows you to force it to not loop with RetryTypes.FALSE.
It uses the fact that list(range(0)) == [], this means that the for loop will never run, however the else will.

It may be better to force self.retry_times to handle if the function will retry, rather than be a state of the module, RetryTypes.{TRUE|FALSE}. Also passing 1 to the function, via self.retry_times, makes it work the same way as passing 0 or less. (Even without my changes). And so it would be confusing if the end user wants to retry once, and the function executes once.
Removing the self.retry_times -= 1 statement removes this problem, as then if I want to retry once if my initial try fails I can.

You can merge RetryTypes.FALSE and RetryTypes.TRUE so your function is less confusing, by retrying if retry_times is set. Then if you want it to save errors pass RetryTypes.REMEMBER_EXCEPTION.

---

If I were to change your code to get all the 'fixes'.

RETRY_TYPES = enum("DEFAULT", "REMEMBER_EXCEPTION")

class Memoizer:
    def __init__(self,
                 retry=RETRY_TYPES.DEFAULT,
                 retry_times=0,
                 exception_type=Exception):
        self.retry = retry
        self.retry_times = retry_times
        self.exception_type = exception_type

    def __call__(self, function):
        d = {}

        def plain_else(args):
            return function(*args)

        def exception_else(args):
            try:
                return function(*args)
            except Exception as e:
                return e

        def plain_return(r_value):
           return r_value

        def exception_return(r_value):
            if isinstance(r_value, Exception):
                raise r_value
            return r_value

        if self.retry is RETRY_TYPES.REMEMBER_EXCEPTION:
            wrapper_else = exception_else
            wrapper_return = exception_return
        else:
            wrapper_else = plain_else
            wrapper_return = plain_return

        @functools.wraps(function)
        def wrapper(*args, overwrite=False):
            if args not in d or overwrite:
                for _ in range(self.retry_times):
                    try:
                        tmp = function(*args)
                    except self.exception_type:
                        continue
                    else:
                        break
                else:
                    tmp = wrapper_else(args)
                d[args] = tmp
            return wrapper_return(d[args])

        return wrapper

It is simpler, and is easier to use / understand.

• If you want to retry pass the amount of times you want to retry.

• If you want to overwrite the memoized variable, pass overwrite=True.
  Forget is ambiguous to what you are forgetting.

  If you want to handle **kwargs at a later date, you may want to change it to overwrite_memoized, or an alternate less common name.

• Less code, and duplication of code.
• Proper retry_times, rather than max_trys. max_trys = retry_times + 1, the initial try is not a retry.
• You only handle one exception, and so exception_types was renamed to exception_type.
• Use a tmp variable, as dict set item is \$O(n)\$ worst case.
  • Prevents the function being \$O(n^2)\$.
  • Follows PEP8 on simple and small try statements.

Answer 2

I don't like the fact that an enumeration with three elements has two named TRUE and FALSE. It sound like the enumeration is simply a boolean, and a third element is not expected. I think they should be called NEVER and ALWAYS or something similar. Those names allude to a possibility of more elements.

The RetryType enum is actually not intuitive at all. When the type is REMEMBER_EXCEPTION or TRUE, retry_times controls whether or not to do retries. When the type is FALSE the retry_times attribute doesn't do anything. This is what lead to my confusion. I suggest you remove RetryTypes and separate it's meaning to two arguments/attributes: retry_times (zero means to not retry) and remember_exception which is a boolean.

If you make these changes __call__ could be simplified a lot. You can define just one wrapper function that handles all cases. It just needs an if self.remember_exception statement, the retry loop should just be skipped if self.retry_times is less than one.

Document your code! I assumed that REMEMBER_EXCEPTION was only meant for use with retry turned on. It was not obvious that it's also designed to be used when retry_times is zero. Use docstrings at the top of modules, classes, methods and functions to make use of your code clear.

Don't import a module inside a nested function somewhere in a file (like in test_retry_cache_errors_eventual_success). Import all modules at the top of your files to be consistent and to provide an overview of what modules a given file makes use of.

Add a test that checks the return value of a decorated function against the return value of the same function not decorated.

def f():
    # ...

df = Memoizer()(f)
self.assertIs(f(), df())