pickle-debug

Pickle expensive operations for debugging operation

When debugging code that runs after a slow step (loading a large file, a long computation, a network call), re-running the slow step on every iteration wastes time. pickle_debug caches the return value of selected functions on disk, so the first call computes and stores the result and later calls with the same arguments load it from the cache.

Requires Python 3.12 or later. No dependencies.

Usage

import pickle_debug

@pickle_debug.cached
def load_data(path, limit=None):
    ...  # slow

pickle_debug.enable("/tmp/mycache")

load_data("big.csv")   # computed, then stored
load_data("big.csv")   # read from cache

Caching is off until enable() is called. Until then, decorated functions simply call through. This makes it safe to leave the @cached decorators in place and turn caching on only while debugging.

Decorating a function from another module

A decorator is an ordinary function, so you can apply it after importing:

import other_module
other_module.slow = pickle_debug.cached(other_module.slow)

This affects callers that look up other_module.slow at call time. Modules that already did from other_module import slow keep the undecorated function, so patch before they import it.

Methods

slow = Slow()
slow.invoke = pickle_debug.cached(slow.invoke)   # one instance

Slow.invoke = pickle_debug.cached(Slow.invoke)   # all instances

The decorator also works inside a class body. For a classmethod, put @cached below @classmethod:

class Slow:
    @pickle_debug.cached
    def invoke(self, x):
        ...

    @classmethod
    @pickle_debug.cached
    def build(cls, x):
        ...

A wrapper may be needed to cache some functions, e.g. a method returning a generator:

def something(x,y,z)->Generator:

@pickle_debug.cached
def cache_something(x,y,z)->list:
	return list(something(x,y,z))

When a cached function is called as a method, self (or cls) is keyed by its class's name (module.QualifiedName), not its current value. All instances of a class share cache entries, and the instance does not need to be picklable. Subclasses get their own entries. Static methods are keyed like plain functions.

API

enable(directory)

Turn caching on, storing results in directory (a str or Path). The directory is created if needed and checked for writability. Raises OSError if it cannot be created or written.

cached(func)

Decorator. Returns a wrapper (using functools.wraps) that:

1. calls func directly if caching is not enabled;
2. otherwise builds a key from func.__module__, func.__qualname__, the positional arguments and the sorted keyword arguments, and hashes it with SHA-256. If func is called as a method, the first argument (self or cls) is replaced by its class's name (see Methods);
3. returns the pickled result from <directory>/<hash>.pickle if present;
4. otherwise calls func, writes the result to the cache, and returns it.

Results are written to a temporary file and then renamed, so an interrupted run does not leave a partial cache entry.

Unpicklable attributes

If an object (in the result or in the arguments) has attributes that cannot be pickled, such as locks, open files, sockets or lambdas, those attributes are left out and the rest of the object is cached. This applies to nested objects too, and covers both __dict__ and __slots__ attributes. Each dropped attribute is logged at DEBUG level.

An object loaded from the cache does not have the dropped attributes, so accessing them raises AttributeError. If the code you are debugging needs them, recreate them after the call.

Only attributes are filtered. An unpicklable value at the top level, or directly inside a list, tuple, dict or set, still raises an exception. So does an object with a custom __reduce__ or a __getstate__ that does not return a dict.

reset()

Delete the whole cache directory. It refuses (raises RuntimeError) if the directory contains anything other than .pickle and .tmp files, to avoid deleting unrelated data if enable() was pointed at the wrong place. It does nothing if caching is not enabled or the directory does not exist.

Caveats

• Arguments and results must be picklable, apart from object attributes, which are left out (see Unpicklable attributes). The arguments are pickled to build the key, and the result is pickled to store it. Because dropped attributes are not part of the key, two arguments that differ only in unpicklable attributes share a cache entry.
• Wrappers may be required for methods that return generators.
• The key is the pickled arguments. Equal arguments that pickle differently (for example sets with a different iteration order) may miss the cache. Objects whose pickle changes between calls never hit it.
• The key does not include the function's code or any files it reads. If you change the function or its input data, call reset() or use a new cache directory, or you will get stale results.
• Instance state is not part of the key. Methods are keyed by the class of self or cls, not by its value, so all instances of a class share entries. If a method's result depends on instance attributes, instances with different state will get each other's cached results. Pass that state as an argument instead, or call reset() between runs.
• This is a debugging aid. Do not use it as a production cache. Cached files are loaded with pickle, so only use a directory you trust.

Logging

The module logs to the pickle_debug logger at DEBUG level when it enables caching, reads from the cache, stores a result, or leaves an unpicklable attribute out.