Sheerka-Old/src/sheerkapickle/handlers.py

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__)