Search code examples
pythonpython-3.xdispatchlrusingle-dispatch

lru_cache interferes with type checking done by single_dispatch


I have a method dispatch decorator with three registered functions. One dispatches on int, which works fine. The second dispatched on a custom type, also works fine. The third is also a custom type, but the Class is wrapped with the lru_cache decorator.

(To make things a little more complicated, the class is instantiated in a roundabout way via a methoddispatch on the __call__ method of another class.)

@lru_cache(maxsize=None, typed=True)
class QualifiedInterval:
    # stuff that works

Inside the Pitch class:

@oph_utils.method_dispatch
def augmented(self, other):
    raise NotImplementedError

@augmented.register(int)
def _(self, other):
    return "works fine"


@augmented.register(Interval)
def _(self, other):
    return "works fine too"

@augmented.register(QualifiedInterval)
def _(self, other):
    return "only works if QualifiedInterval class does NOT have lru_cache"

(There's a lot more going on, but this is the bits that don't work.)

Basically - if I have lru_cache, and pass a QualifiedInterval into the function, it does not dispatch and raises NotImplementedError. If I comment out the cache decorator, it works. And manual type checking at the REPL shows the same type ("QualifiedInterval") either way. I've tried calling the command that created the QualifiedInterval several different ways, and tried assigning it to a variable. Still doesn't work. I've tried doing explicit typechecking in the Augmented function. The typecheck fails there as well, if caching is enabled.


Solution

  • Analysis of What is Going Wrong

    Basically - if I have lru_cache, and pass a QualifiedInterval into the function, it does not dispatch

    The lru_cache is function that returns a decorator that wraps any callable (including classes). So when you apply the lru_cache to the QualifiedInterval class, that variable becomes assigned to the wrapper function rather than the class itself.

    >>> @lru_cache(maxsize=None, typed=True)
    class QualifiedInterval:
        pass
    
    >>> type(QualifiedInterval)
    <class 'functools._lru_cache_wrapper'>
    

    Single dispatch works by matching the type of the first parameter to the appropriate method. However, when you pass in an instance of QualifiedInterval, its type doesn't match functools._lru_cache_wrapper, so single dispatch falls back to the base method(which raises NotImplemented.

    Solution

    Teach single dispatch to match on the actual original class (type) instead of the wrapped class:

    @augmented.register(QualifiedInterval.__wrapped__)
    def _(self, other):
        return "show now work QualifiedInterval class has an lru_cache"
    

    Note the addition of the .__wrapped__ attribute which reaches through the wrapper function to get to the original unwrapped class.

    Hope that clears it all up and shows the way forward :-)