Memoize a function so that it isn't reset when I rerun the file in Python

Question

I often do interactive work in Python that involves some expensive operations that I don't want to repeat often. I'm generally running whatever Python file I'm working on frequently.

If I write:

import functools32

@functools32.lru_cache()
def square(x):
    print "Squaring", x
    return x*x

I get this behavior:

>>> square(10)
Squaring 10
100
>>> square(10)
100
>>> runfile(...)
>>> square(10)
Squaring 10
100

That is, rerunning the file clears the cache. This works:

try:
    safe_square
except NameError:
    @functools32.lru_cache()
    def safe_square(x):
        print "Squaring", x
        return x*x

but when the function is long it feels strange to have its definition inside a try block. I can do this instead:

def _square(x):
    print "Squaring", x
    return x*x

try:
    safe_square_2
except NameError:
    safe_square_2 = functools32.lru_cache()(_square)

but it feels pretty contrived (for example, in calling the decorator without an '@' sign)

Is there a simple way to handle this, something like:

@non_resetting_lru_cache()
def square(x):
    print "Squaring", x
    return x*x

?

Answer 1

There is no persistent_lru_cache in the stdlib. But you can build one pretty easily.

The functools source is linked directly from the docs, because this is one of those modules that's as useful as sample code as it is for using it directly.

As you can see, the cache is just a dict. If you replace that with, say, a shelf, it will become persistent automatically:

def persistent_lru_cache(filename, maxsize=128, typed=False):
    """new docstring explaining what dbpath does"""
    # same code as before up to here
    def decorating_function(user_function):
        cache = shelve.open(filename)
        # same code as before from here on.

Of course that only works if your arguments are strings. And it could be a little slow.

So, you might want to instead keep it as an in-memory dict, and just write code that pickles it to a file atexit, and restores it from a file if present at startup:

    def decorating_function(user_function):
        # ...

        try:
            with open(filename, 'rb') as f:
                cache = pickle.load(f)
            except:
                cache = {}
        def cache_save():
            with lock:
                with open(filename, 'wb') as f:
                    pickle.dump(cache, f)
        atexit.register(cache_save)

        # …
        wrapper.cache_save = cache_save
        wrapper.cache_filename = filename

Or, if you want it to write every N new values (so you don't lose the whole cache on, say, an _exit or a segfault or someone pulling the cord), add this to the second and third versions of wrapper, right after the misses += 1:

            if misses % N == 0:
                cache_save()

---

See here for a working version of everything up to this point (using save_every as the "N" argument, and defaulting to 1, which you probably don't want in real life).

If you want to be really clever, maybe copy the cache and save that in a background thread.

You might want to extend the cache_info to include something like number of cache writes, number of misses since last cache write, number of entries in the cache at startup, …

And there are probably other ways to improve this.

From a quick test, with save_every=1, this makes the cache on both get_pep and fib (from the functools docs) persistent, with no measurable slowdown to get_pep and a very small slowdown to fib the first time (note that fib(100) has 100097 hits vs. 101 misses…), and of course a large speedup to get_pep (but not fib) when you re-run it. So, just what you'd expect.