import datetime import re import threading import uuid from sheerkapickle import utils class ToReduce: def __init__(self, predicate, get_value): self.predicate = predicate self.get_value = get_value class SheerkaRegistry(object): def __init__(self): self._handlers = {} self._base_handlers = {} def get(self, cls_or_name, default=None): """ :param cls_or_name: the type or its fully qualified name :param default: default value, if a matching handler is not found Looks up a handler by type reference or its fully qualified name. If a direct match is not found, the search is performed over all handlers registered with base=True. """ handler = self._handlers.get(cls_or_name) # attempt to find a base class if handler is None and utils.is_type(cls_or_name): for cls, base_handler in self._base_handlers.items(): if issubclass(cls_or_name, cls): return base_handler return default if handler is None else handler def register(self, cls, handler=None, base=False): """Register the a custom handler for a class :param cls: The custom object class to handle :param handler: The custom handler class (if None, a decorator wrapper is returned) :param base: Indicates whether the handler should be registered for all subclasses This function can be also used as a decorator by omitting the `handler` argument:: @jsonpickle.handlers.register(Foo, base=True) class FooHandler(jsonpickle.handlers.BaseHandler): pass """ if handler is None: def _register(handler_cls): self.register(cls, handler=handler_cls, base=base) return handler_cls return _register if not utils.is_type(cls): raise TypeError('{!r} is not a class/type'.format(cls)) # store both the name and the actual type for the ugly cases like # _sre.SRE_Pattern that cannot be loaded back directly self._handlers[utils.importable_name(cls)] = \ self._handlers[cls] = handler if base: # only store the actual type for subclass checking self._base_handlers[cls] = handler def unregister(self, cls): self._handlers.pop(cls, None) self._handlers.pop(utils.importable_name(cls), None) self._base_handlers.pop(cls, None) registry = SheerkaRegistry() register = registry.register unregister = registry.unregister get = registry.get class BaseHandler(object): def __init__(self, sheerka, context): """ Initialize a new handler to handle a registered type. :Parameters: - `context`: reference to pickler/unpickler """ self.sheerka = sheerka self.context = context def __call__(self, sheerka, context): """This permits registering either Handler instances or classes :Parameters: - `context`: reference to pickler/unpickler """ self.sheerka = sheerka self.context = context return self def flatten(self, obj, data): """ Flatten `obj` into a json-friendly form and write result to `data`. :param object obj: The object to be serialized. :param dict data: A partially filled dictionary which will contain the json-friendly representation of `obj` once this method has finished. """ raise NotImplementedError('You must implement flatten() in %s' % self.__class__) def new(self, data): raise NotImplementedError('You must implement new() in %s' % self.__class__) def restore(self, data, instance): """ Restore an object of the registered type from the json-friendly representation `obj` and return it. """ raise NotImplementedError('You must implement restore() in %s' % self.__class__) @classmethod def handles(self, cls): """ Register this handler for the given class. Suitable as a decorator, e.g.:: @MyCustomHandler.handles class MyCustomClass: def __reduce__(self): ... """ registry.register(cls, self) return cls # class DatetimeHandler(BaseHandler): # """Custom handler for datetime objects # # Datetime objects use __reduce__, and they generate binary strings encoding # the payload. This handler encodes that payload to reconstruct the # object. # # """ # # def flatten(self, obj, data): # pickler = self.context # if not pickler.unpicklable: # return str(obj) # cls, args = obj.__reduce__() # flatten = pickler.flatten # payload = utils.b64encode(args[0]) # args = [payload] + [flatten(i, reset=False) for i in args[1:]] # data['__reduce__'] = (flatten(cls, reset=False), args) # return data # # def restore(self, data): # cls, args = data['__reduce__'] # unpickler = self.context # restore = unpickler.restore # cls = restore(cls, reset=False) # value = utils.b64decode(args[0]) # params = (value,) + tuple([restore(i, reset=False) for i in args[1:]]) # return cls.__new__(cls, *params) # # # DatetimeHandler.handles(datetime.datetime) # DatetimeHandler.handles(datetime.date) # DatetimeHandler.handles(datetime.time) class RegexHandler(BaseHandler): """Flatten _sre.SRE_Pattern (compiled regex) objects""" def flatten(self, obj, data): data['pattern'] = obj.pattern return data def new(self, data): return re.compile(data['pattern']) def restore(self, data, instance): return instance RegexHandler.handles(type(re.compile(''))) class UUIDHandler(BaseHandler): """Serialize uuid.UUID objects""" def flatten(self, obj, data): data['hex'] = obj.hex return data def new(self, data): return uuid.UUID(data['hex']) def restore(self, data, instance): return instance UUIDHandler.handles(uuid.UUID) class LockHandler(BaseHandler): """Serialize threading.Lock objects""" def flatten(self, obj, data): data['locked'] = obj.locked() return data def new(self, data): lock = threading.Lock() if data.get('locked', False): lock.acquire() return lock def restore(self, data, instance): return instance _lock = threading.Lock() LockHandler.handles(_lock.__class__)