Spaces:
Running
Running
| from functools import lru_cache, singledispatch | |
| from typing import Any, Callable, List, Tuple, Union | |
| import attr | |
| class _DispatchNotFound: | |
| """A dummy object to help signify a dispatch not found.""" | |
| pass | |
| class MultiStrategyDispatch: | |
| """ | |
| MultiStrategyDispatch uses a combination of exact-match dispatch, | |
| singledispatch, and FunctionDispatch. | |
| """ | |
| __slots__ = ( | |
| "_direct_dispatch", | |
| "_function_dispatch", | |
| "_single_dispatch", | |
| "_generators", | |
| "dispatch", | |
| ) | |
| def __init__(self, fallback_func): | |
| self._direct_dispatch = {} | |
| self._function_dispatch = FunctionDispatch() | |
| self._function_dispatch.register(lambda _: True, fallback_func) | |
| self._single_dispatch = singledispatch(_DispatchNotFound) | |
| self.dispatch = lru_cache(maxsize=None)(self._dispatch) | |
| def _dispatch(self, cl): | |
| try: | |
| dispatch = self._single_dispatch.dispatch(cl) | |
| if dispatch is not _DispatchNotFound: | |
| return dispatch | |
| except Exception: | |
| pass | |
| direct_dispatch = self._direct_dispatch.get(cl) | |
| if direct_dispatch is not None: | |
| return direct_dispatch | |
| return self._function_dispatch.dispatch(cl) | |
| def register_cls_list(self, cls_and_handler, direct: bool = False): | |
| """ register a class to direct or singledispatch """ | |
| for cls, handler in cls_and_handler: | |
| if direct: | |
| self._direct_dispatch[cls] = handler | |
| else: | |
| self._single_dispatch.register(cls, handler) | |
| self.clear_direct() | |
| self.dispatch.cache_clear() | |
| def register_func_list( | |
| self, | |
| func_and_handler: List[ | |
| Union[ | |
| Tuple[Callable[[Any], bool], Any], | |
| Tuple[Callable[[Any], bool], Any, bool], | |
| ] | |
| ], | |
| ): | |
| """register a function to determine if the handle | |
| should be used for the type | |
| """ | |
| for tup in func_and_handler: | |
| if len(tup) == 2: | |
| func, handler = tup | |
| self._function_dispatch.register(func, handler) | |
| else: | |
| func, handler, is_gen = tup | |
| self._function_dispatch.register( | |
| func, handler, is_generator=is_gen | |
| ) | |
| self.clear_direct() | |
| self.dispatch.cache_clear() | |
| def clear_direct(self): | |
| """Clear the direct dispatch.""" | |
| self._direct_dispatch.clear() | |
| class FunctionDispatch: | |
| """ | |
| FunctionDispatch is similar to functools.singledispatch, but | |
| instead dispatches based on functions that take the type of the | |
| first argument in the method, and return True or False. | |
| objects that help determine dispatch should be instantiated objects. | |
| """ | |
| __slots__ = ("_handler_pairs",) | |
| def __init__(self): | |
| self._handler_pairs = [] | |
| def register( | |
| self, can_handle: Callable[[Any], bool], func, is_generator=False | |
| ): | |
| self._handler_pairs.insert(0, (can_handle, func, is_generator)) | |
| def dispatch(self, typ): | |
| """ | |
| returns the appropriate handler, for the object passed. | |
| """ | |
| for can_handle, handler, is_generator in self._handler_pairs: | |
| # can handle could raise an exception here | |
| # such as issubclass being called on an instance. | |
| # it's easier to just ignore that case. | |
| try: | |
| ch = can_handle(typ) | |
| except Exception: | |
| continue | |
| if ch: | |
| if is_generator: | |
| return handler(typ) | |
| else: | |
| return handler | |
| raise KeyError("unable to find handler for {0}".format(typ)) | |