A module wrapping the python module `klepto`. Function names are considered by the keygen and the decorators of caches can be applied to classes.

wrap_klepto.py

import functools
import importlib
import inspect
import sys
from copy import copy
from pathlib import Path
from types import ModuleType

import klepto

__all__ = list(attr for attr in dir(klepto) if attr[0] != '_')

def init_app(app):
    """Do nothing"""
    pass

def _clone_klepto_module(module):
    submodule_name = module.__name__[7:] #  len('klepto.') = 7
    spec = importlib.util.find_spec(module.__name__)
    new_module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(new_module)
    sys.modules[__name__ + '.' + submodule_name] = new_module
    return new_module

_archives = _clone_klepto_module(klepto._archives)
_cache = _clone_klepto_module(klepto._cache)
_inspect = _clone_klepto_module(klepto._inspect)
_pickle = _clone_klepto_module(klepto._pickle)
archives = _clone_klepto_module(klepto.archives)
crypto = _clone_klepto_module(klepto.crypto)
info = _clone_klepto_module(klepto.info)
keymaps = _clone_klepto_module(klepto.keymaps)
rounding = _clone_klepto_module(klepto.rounding)
safe = _clone_klepto_module(klepto.safe)
tools = _clone_klepto_module(klepto.tools)

def _wrap_keygen(_keygen):
    # `inspect` will be deleted at the end of this module, so make references
    # of the methods will be used in wrapped here.
    inspect_isclass = inspect.isclass
    inspect_signature = inspect.signature
    def _wrapped_keygen(func, ignore, *args, **kwargs):
        kwargs['__func_name_cache__'] = func.__name__

        if hasattr(func, '__wrapped__') and not hasattr(func, '__signature__'):
            # `inspect.signature` follows wrapped functions by default
            # The arguments may change in the wrapping chain. It may be hard to
            # use this decorator if it handle the maximum level of the wrapped
            # functions that `inspect.signature` will follow as an argument.
            func.__signature__ = inspect_signature(func)

        # patch for the issue: https://github.com/uqfoundation/klepto/issues/75
        if inspect_isclass(func):
            # magic key used in the wrapper for the caches
            kwargs.pop('__to_original_m061p3__', None)
            ignore = set(ignore)
            args = (None,) + args
            # Magic! Explain:
            # >>> f = getattr(None, '__init__')
            # >>> n = getattr(f, '__self__')
            # >>> n is None
            # True
            while inspect_isclass(func) and func is not type:
                if '__init__' in func.__dict__:
                    func = func.__init__
                    # `__init__` usually use "self" as the first argument
                    ignore.add('self')
                    break
                elif '__new__' in func.__dict__:
                    func = copy(func.__new__)
                    # lie the checker for bounded methods
                    func.__name__ = '__init__'
                    # `__new__` usually use "cls" as the first argument
                    ignore.add('cls')
                    break
                else:
                    func = func.__class__
            func = (lambda: None) if func is type else func
            if hasattr(func, '__wrapped__'):
                func.__signature__ = inspect_signature(func)
        _args, _kwargs = _keygen(func, ignore, *args, **kwargs)
        return _args, _kwargs
    return _wrapped_keygen

# `klepto._cache._keygen` is used by each cache object in `klepto._cache` to 
# generate the arguments passed to the keymaps which encode the input arguments
# to the keys.
_cache._keygen = _wrap_keygen(_cache._keygen)

_CACHES_PATH = Path(__file__).parent / 'caches'
if not _CACHES_PATH.is_dir():
    _CACHES_PATH.mkdir()

# __archives: A dictionary caches the loaded archives to make sure that
#     the different archives getten at different part of the code but
#     getten with the same namespace can access the same memory cache.
__archives = {}

def _decorate_archive(archive):
    @functools.wraps(archive)
    def wrapper(namespace, *args, dict=None, cached=True, **kwargs):
        """
        Arguments:
            namespace: A string of the namespace which is the prefix of the file
                name or the prefix of the folder name of the archive. The diffe-
                rent archive instance with the same namespace link to the same 
                directory or the same file.
            dict: Initial dictionary to seed the archive. This only affect the
                memory part of the archive.
            cached: If True, use an in-memory cache interface to the archive.
        """
        archive_type = archive.__name__[:-8]
        name = f'{namespace}-{archive_type}'
        path = _CACHES_PATH / name
        if name not in __archives:
            __archives[name] = archive(
                *args, name=path, dict=None, cached=True, **kwargs)
        return __archives[name]
    return wrapper

# wrap archives in `klepto.archives`
def _wrap_archives():
    for attr in archives.__all__:
        if attr.endswith('_archive'):
            original_archive = getattr(archives, attr)
            new_archive = _decorate_archive(original_archive)
            setattr(archives, attr, new_archive)
_wrap_archives()

# assign non-module variable
def _link_klepto_attr(attr):
    instance = getattr(klepto, attr)
    # e.g. instance.__module__ = klepto._cache
    module_chain = instance.__module__.split('.')[1:]
    if not module_chain:
        raise ValueError(
            'Should not link the attributes not defined ' \
            'in the submodules of `klepto`')
    module = globals()[module_chain[0]]
    for name in module_chain[1:]:
        module = getattr(module, name)
    return getattr(module, attr)

no_cache = _link_klepto_attr('no_cache')
inf_cache = _link_klepto_attr('inf_cache')
lfu_cache = _link_klepto_attr('lfu_cache')
lru_cache = _link_klepto_attr('lru_cache')
mru_cache = _link_klepto_attr('mru_cache')
rr_cache = _link_klepto_attr('rr_cache')
signature = _link_klepto_attr('signature')
isvalid = _link_klepto_attr('isvalid')
validate = _link_klepto_attr('validate')
keygen = _link_klepto_attr('keygen')
strip_markup = _link_klepto_attr('strip_markup')
NULL = _link_klepto_attr('NULL')
_keygen = _link_klepto_attr('_keygen')

def _wrap_cache(cache):
    """Factory for making original caches support classes

    Because `pickle` cannot unpickle the object of which class cannot be 
    imported directly by the full name of the class, see the first refrence
    below for the source code of cpython, it will fail when unpickling the 
    classes wrapped by the original caches. To make the caches support classes,
    the original classes are returned by the new caches, as the wrappers, but
    `__new__` and `__init__` of the classes are hooked, and the wrapped function
    is set as an attribute named `__cache_instances_factory__` of the class.

    Returns:
        Original cache class which `__call__` is be replaced.

    Notices:
        1. Not support the classes whose metaclass are not `type`.
        2. Not influent the subclasses.

    Related References:
        https://github.com/python/cpython/blob/0378d98678f3617fd44d9a6266e7c17ebce62755/Lib/pickle.py#L1063
        https://github.com/uqfoundation/klepto/issues/75
        https://docs.python.org/3.7/library/pickle.html#pickling-class-instances
    """
    # `functools` and `inspect` will be deleted at the end of this module.
    functools_wraps = functools.wraps
    inspect_isclass = inspect.isclass
    inspect_getfullargspec = inspect.getfullargspec

    # When `cache.__call__` is called in `wrapper(self, user_function)`,
    # `cache.__call__` is replaced, so it need to be cached here.
    original_call = cache.__call__
    @functools.wraps(cache.__call__)
    def wrapper(self, user_function):
        decorated_function = original_call(self, user_function)
        if not inspect_isclass(user_function):
            return decorated_function
        else:
            user_class = user_function # alias
            user_class.__cache_instances_factory__ = decorated_function

            # Use the original `__getnewargs_ex__` or `__getnewargs__` if they
            # are exists.
            if not hasattr(user_class, '__getnewargs_ex__') and \
                not hasattr(user_class, '__getnewargs__'):
                # XXX: If the class is wrapped by `xx_cache` more than once,
                #   this will get the correct argspec. Should the situation
                #   of wrapping the class more than once be supported?
                argspec = inspect_getfullargspec(user_class.__new__)
                def __getnewargs_ex__(self):
                    # `__newargs__` and `__newkwargs__` are set in decorated
                    # `__new__`.
                    if argspec.varargs is not None:
                        args = self.__newargs__
                    else:
                        args = self.__newargs__[:len(argspec.args)-1]
                    if argspec.varkw is not None:
                        kwargs = self.__newkwargs__
                    else:
                        kwonlyargs = set(argspec.kwonlyargs)
                        kwonlyargs.add('__to_original_m061p3__')
                        kwargs = {k: self.__newkwargs__[k] 
                            for k in kwonlyargs if k in self.__newkwargs__}
                    return args, kwargs
                user_class.__getnewargs_ex__ = __getnewargs_ex__
            
            # fix inheritance mechanism
            if '__new__' in user_class.__dict__ or \
                not hasattr(user_class, '__original_new__'):
                user_class.__original_new__ = original_new = user_class.__new__
            else:
                # give the subclass the original `__new__`
                original_new = user_class.__original_new__
            @functools_wraps(user_class.__new__)
            def wrap_new(cls, *args, **kwargs):
                # If __to_original_m061p3__ is in `kwargs`, it means that this
                # function call comes from `__cache_instances_factory__` or that
                # `__cache_instances_factory__` should be bypassed.
                # only pass to the factory if the variable, `cls`, is itself
                if '__to_original_m061p3__' not in kwargs and cls is user_class:
                    kwargs['__to_original_m061p3__'] = True
                    return cls.__cache_instances_factory__(*args, **kwargs)
                else:
                    kwargs.pop('__to_original_m061p3__', None)
                    if original_new is object.__new__:
                        new_obj = original_new(cls)
                    else:
                        new_obj = original_new(cls, *args, **kwargs)
                    # For `__getnewargs_ex__`.
                    # When unpickling the object, `__cache_instances_factory__`
                    # should be bypassed.
                    kwargs['__to_original_m061p3__'] = True
                    new_obj.__newargs__, new_obj.__newkwargs__ = args, kwargs
                    return new_obj
            user_class.__new__ = wrap_new 

            original_init = user_class.__init__
            @functools_wraps(user_class.__init__)
            def wrap_init(self, *args, **kwargs):
                if '__to_original_m061p3__' in kwargs:
                    kwargs.pop('__to_original_m061p3__')
                    original_init(self, *args, **kwargs)
                else:
                    # If the instance is not created by 
                    # `__cache_instances_factory__`, then do nothing.
                    #
                    # If the warpped class is a direct subclass of `object`
                    # and of which metaclass is `type`, the proccess of
                    # creating an instance as the following.
                    # 
                    # decorated __new__ without `__to_original_m061p3__` ->
                    # __cache_instances_factory__ ->
                    # decorated __new__ with `__to_original_m061p3__` ->
                    # original __new__ without `__to_original_m061p3__` ->
                    # decorated __init__ without `__to_original_m061p3__` ->
                    # decorated __init__ with `__to_original_m061p3__`
                    #
                    # For more details about how instances are constructed,
                    # refer to:
                    # https://blog.ionelmc.ro/2015/02/09/understanding-python-metaclasses/#putting-it-all-together
                    pass
            user_class.__init__ = wrap_init
            return user_class
    cache.__call__ = wrapper
    return cache

no_cache = _wrap_cache(no_cache)
inf_cache = _wrap_cache(inf_cache)
lfu_cache = _wrap_cache(lfu_cache)
lru_cache = _wrap_cache(lru_cache)
mru_cache = _wrap_cache(mru_cache)
rr_cache = _wrap_cache(rr_cache)

del _clone_klepto_module
del _decorate_archive
del _link_klepto_attr
del _wrap_archives
del _wrap_cache
del _wrap_keygen
del copy
del functools
del importlib
del inspect
del klepto
del ModuleType
del Path
del sys