python-preserve/preserve.py

"""
Preserve, for when pickling is just a bit too intense

Preserve is a simple framework to painlessly convert your python objects and
data structures into a nested structure containing only primitive types.

While intended for use in serialisation, Preserve stops short of the actual
serialisation bit - you can then use whatever method you prefer to dump
the nested primitive structure to JSON, TOML, or some other message packing
or storage system - most of which handle all primitive types by default.

"""

# pylint: disable=bad-staticmethod-argument
import inspect
from functools import partial
from collections import namedtuple
from types import SimpleNamespace, MappingProxyType


class RestoreMethod():
    """
    Contains the various supplied restore method implementations. All methods here are intended
    to be passed in when creating a PreservableType.

    All methods must use the same argument set: `(cls, state)`, where `cls` is the class usually
    stored as the `type` field in the Preservable, and `state` is whatever was previously
    returned by `preserve()`. The value returned by the method is intended to be a new copy of the
    instance originally passed to `preserve()`.
    """
    @staticmethod
    def default_restore(cls, state):
        """
        The default restore method used by Preserve. Creates a new instance of the class (without
        calling `__init__()`), and directly updates the object `__dict__` with the state passed
        in.

        If the object has implemented a `__restore_init__()` method, that is then called.
        """
        obj = cls.__new__(cls)
        obj.__dict__.update(state)
        if inspect.ismethod(getattr(obj, "__restore_init__", None)):
            obj.__restore_init__()
        return obj

    @staticmethod
    def restore_after_init(cls, state):
        """
        Does the same things as the default restore method, but calls `__init__()` on the new
        object first, without any arguments. Useful when you want to retain existing init
        behaviour but override some of the object attributes immediately afterward from the
        preserved state. Often used along with `include_attrs` when you don't want to
        implement `__restore_init__()`.

        Will not call `__restore_init__()` on the object, even if it is implemented.
        """
        obj = cls()
        obj.__dict__.update(state)
        return obj

    @staticmethod
    def setstate(cls, state):
        """
        Uses the `__setstate__(self, state)` method on the class - the same protocol used for
        things like Pickle.
        """
        obj = cls.__new__(cls)
        obj.__setstate__(state)
        return obj

    @staticmethod
    def pass_to_new(cls, state):
        """
        Creates the new instance by directly passing `state` to the `__new__` method on `cls`.
        Useful for immutable types (like tuples) that are initialised this way.
        """
        obj = cls.__new__(cls, state)
        return obj


class RestoreMethodGenerator():
    """
    Similar to RestoreMethod, but contains functions intended to be _called_ that will _return_ a
    RestoreMethod. These will usually be a partial function with a given set of arguments to
    allow customisation of the restore method.
    """

    @staticmethod
    def restore_after_init_args(*args, **keywords):
        """
        Returns a PreserveMethod similar to `restore_after_init()`, but allows you to pass in
        a set of args and kwargs that will be passed on to the `__init__()` call.
        """
        def restore_after_init_args(cls, state, *args, **keywords):
            obj = cls(*args, **keywords)
            obj.__dict__.update(state)
            return obj
        return partial(restore_after_init_args, *args, **keywords)


class PreserveMethod():
    """
    Contains the various supplied preserve method implementations. Methods here are intended
    to be passed directly in when creating a PreservableType.

    All methods must only accept a single argument, `obj`: the object that needs to be preserved.
    """

    @staticmethod
    def default_preserve(obj):
        """
        The default preserve method used by Preserve. Simply grabs the `__dict__` attribute of the
        object (returned by `vars(obj)`). Will only work if all object attributes are themselves
        preservable.
        """
        return vars(obj)

    @staticmethod
    def getstate(obj):
        """
        Uses the `__getstate__(obj)` _method_ on the class - the same protocol used for
        things like Pickle.
        """
        return obj.__getstate__()


class PreserveMethodGenerator():
    """
    Similar to PreserveMethod, but contains functions intended to be _called_ that will _return_ a
    PreserveMethod. These will usually be a partial function with a given set of arguments to
    allow customisation of the preserve method.
    """

    @staticmethod
    def default_preserve_filter(include_attrs=None, exclude_attrs=None):
        """
        Returns a PreserveMethod similar to `default_preserve()` that will filter out certain
        attributes before returning them in the `state` dictionary. Both arguments should be
        iterables of argument names (strings).

        Passing None to `include_attrs` will skip the inclusion filter.
        """
        def default_preserve_filter(obj, include_attrs, exclude_attrs):
            state = vars(obj)
            attr_keys = state.keys()
            if include_attrs:
                include_attrs = set(include_attrs)
                attr_keys = attr_keys & include_attrs
            if exclude_attrs:
                exclude_attrs = set(exclude_attrs)
                attr_keys = attr_keys - exclude_attrs

            return {key: state[key] for key in attr_keys}

        return partial(default_preserve_filter, include_attrs=include_attrs,
                       exclude_attrs=exclude_attrs)

    @staticmethod
    def coerce(new_type):
        """
        Returns a PreserveMethod that simply passes the object into `new_type` to try and
        create a new object of that type. Useful for simplifying objects when preserving them.
        """
        def coerce(obj, new_type):
            return new_type(obj)
        return partial(coerce, new_type=new_type)


rm = RestoreMethod
rm_gen = RestoreMethodGenerator
pm = PreserveMethod
pm_gen = PreserveMethodGenerator


escape_key = "<_jam>"
"""
The string used as a dict key in the preserved output to indicate Preserve metadata.

Set to <_jam>` by default.
"""

escaped_state_key = "<_jam_state>"
"""
The string used as a dict key in the preserved output to indicate Preserve object state.

Set to <_jam_state>` by default.
"""

PreservableType = namedtuple('PreservableRecord', ['type',
                                                   'name',
                                                   'preserve_method',
                                                   'restore_method'])

common = SimpleNamespace()
common.tuple = PreservableType(tuple, 'tuple', pm_gen.coerce(list), rm.pass_to_new)


_preserve_types = {}
"""
An internal dict of PreservableTypes that can be preserved. Keys are the type.
"""

_restore_types = {}
"""
An internal dict of PreservableTypes that can be restored. Keys are the names.
"""

preservables = MappingProxyType(_restore_types)
"""
A dict of all PreservableTypes Preserve is currently able to handle, with the names as keys.
"""

raw_preservables = [dict, list, str, int, float, bool, type(None)]
"""
A list of types that Preserve will pass directly through to the preserved output. Note that both
`dict` and `list` are special cases, must always be present in this list (attempting to remove
them will break things).
"""


def register(preservable_type):
    """
    Registers an existing PerservableType for use with Preserve. Intended to be used with provided
    PreservableTypes from somewhere like preserve.common.

    To add a custom preservable class, instead consider `add_preservable()` or
    the `@preservable` class decorator.

    Args:
        preservable_type: An instance of preserve.PreservableType
    """
    _preserve_types[preservable_type.type] = preservable_type
    _restore_types[preservable_type.name] = preservable_type


def add_preservable(preservable_type, name, preserve_method, restore_method):
    """
    Creates a new PreservableType with the details given and registers it with Preserve.
    """
    if preservable_type in _preserve_types:
        raise Exception(F"Preservable type already registered with type {preservable_type}")

    if type(name) != str:
        raise Exception("Preservable name must be a string")

    if name in _restore_types:
        raise Exception(F"Preservable type already registered with name '{name}'")

    register(PreservableType(preservable_type, name, preserve_method, restore_method))


def preservable(cls=None, *, include_attrs=None, exclude_attrs=None, preserve_method=None,
                restore_method=None, name=None):
    """
    Class decorator to register a class as preservable. Creates and adds a corresponding
    PreservableType to Preserve. Usable either raw (`@preservable` has sane defaults) or with
    arguments (`@preservable(exclude_attrs = 'my_attr')`).

    The default preserve method attempts to directly preserve all object attributes, optionally
    filtering them if `include_attrs` or `exclude_attrs` are present. If `preserve_method` is set,
    `include_attrs` and `exclude_attrs` are ignored.

    The default restore method creates a new instance of the class - without calling __init__() -
    and directly restores any preserved object attributes. If the restored object implements it,
    `__restore_init__()` is then called.

    If any arguments to this decorator remain None, it will look for corresponding class
    attributes on the class - this allows subclasses to inherit Preserve behaviour
    where desirable.
    Class attributes: `_preserve_include_attrs`, `_preserve_exclude_attrs`, `_preserve_method`,
    `_restore_method`, `_preserve_name`

    If `__getstate__()` or `__setstate__()` are implemented in the class (the Pickle protocol),
    they will override the default preserve or restore method, respectively.

    Args:
        include_attrs: A list of attribute names to use when preserving the object. Only used if
            the preserve_method is left as default. Leaving this as none includes all available
            attributes.
        exclude_attrs: A list of attribute names to leave out when preserving the object. Only
            used if the preserve_method is left as default.
        preserve_method: A callable to be used to preserve this preservable. Some common ones and
            details on their requirements are provided in `preserve.PreserveMethod`.
            Leave as None for default behaviour.
        restore_method: A callable to be used to restore this preservable. Some common ones and
            details on their requirements are provided in `preserve.RestoreMethod`
            Leave as None for default behaviour.
        name: The string used to indentify this class in the preserved output.
            Must be unique among all preservables. Defaults to the
            class name if left as None.
    """
    # If called with kwargs rather than as a direct decorator, we need to return
    # a decorator that Python will only pass `cls` into:
    if cls is None:
        return partial(preservable, include_attrs=include_attrs,
                       exclude_attrs=exclude_attrs, preserve_method=preserve_method,
                       restore_method=restore_method, name=name)

    if include_attrs is None:
        include_attrs = getattr(cls, '_preserve_include_attrs', None)
    if exclude_attrs is None:
        exclude_attrs = getattr(cls, '_preserve_exclude_attrs', None)
    if preserve_method is None:
        preserve_method = getattr(cls, '_preserve_method', None)
    if restore_method is None:
        restore_method = getattr(cls, '_restore_method', None)
    if name is None:
        name = getattr(cls, '_preserve_name', None)

    if name is None:
        name = cls.__name__

    if preserve_method is None:
        if hasattr(cls, '__getstate__'):
            preserve_method = pm.getstate
        else:
            if include_attrs or exclude_attrs:
                # String iterable gaurds
                if isinstance(include_attrs, str):
                    include_attrs = (include_attrs,)
                if isinstance(exclude_attrs, str):
                    exclude_attrs = (exclude_attrs,)

                preserve_method = pm_gen.default_preserve_filter(include_attrs, exclude_attrs)
            else:
                preserve_method = pm.default_preserve

    if restore_method is None:
        if hasattr(cls, '__setstate__'):
            restore_method = rm.setstate
        else:
            restore_method = rm.default_restore

    add_preservable(cls, name, preserve_method, restore_method)

    return cls


def preserve(target_obj):
    """
    Preserve target_obj, running through its contents recursively.

    Pass the result back to `restore()` to get a copy of the original object and its content.

    Args:
        target_obj: The object to be preserved. This object and all its nested
            contents must be preservable.

    Returns:
        The preserved data structure - a nested structure containing only types
        from `raw_preservables` (by default dict, list, str, int, float, bool, type(None))
    """
    obj_type = type(target_obj)

    if obj_type == dict:
        dict_jam = {}
        for key, val in target_obj.items():
            if type(key) != str:
                raise Exception("Non-string dictionary keys are not preservable")
            if key in (escape_key, escaped_state_key):
                raise Exception(F"Dict key '{key}' is reserved as an internal escape key")
            dict_jam[key] = preserve(val)
        return dict_jam

    elif obj_type == list:
        list_jam = []
        for val in target_obj:
            list_jam.append(preserve(val))
        return list_jam

    elif obj_type in raw_preservables:
        return target_obj

    elif obj_type in _preserve_types:
        return _preserve_preservable(target_obj, _preserve_types[obj_type])
    else:
        raise Exception(F"Object {target_obj} is not preservable")


def _preserve_preservable(target_obj, preservable_type):
    """
    Internal preserve function specifically to deal with PreservableTypes. Result will always be
    at minimum a dict with an escape key to hold metadata.
    """
    obj_state = preserve(preservable_type.preserve_method(target_obj))

    # This dict is here for future use to store other metadata.
    escaped_metadata = {}

    if (type(obj_state) == dict):
        # Collapse state dict
        obj_jam = obj_state
    else:
        obj_jam = {escaped_state_key: obj_state}

    if len(escaped_metadata) == 0:
        # Collapse metadata dict
        escaped_metadata = preservable_type.name
    else:
        escaped_metadata['name'] = preservable_type.name
    obj_jam[escape_key] = escaped_metadata

    return obj_jam


def restore(obj_jam):
    """
    Restore the result of `preserve()` back into its original form.

    Args:
        obj_jam: The data structure to restore, usually the result of a
            `preserve()` call.
    """
    obj_type = type(obj_jam)

    if obj_type == dict:
        return _restore_dict(obj_jam)

    elif obj_type == list:
        restored_list = []
        for val in obj_jam:
            restored_list.append(restore(val))
        return restored_list

    elif obj_type in raw_preservables:
        return obj_jam

    else:
        raise Exception(F"Object {str(obj_jam)} is not restorable")


def _restore_dict(dict_jam):
    restored_dict = {}
    for key, val in dict_jam.items():
        if type(key) != str:
            raise Exception("Non-string dictionary keys are not restorable")
        if key not in (escape_key, escaped_state_key):
            restored_dict[key] = restore(val)

    if escape_key in dict_jam:
        metadata = dict_jam[escape_key]
        if type(metadata) == dict:
            name = metadata['name']
        else:
            # Metadata was collapsed
            name = metadata

        if name not in _restore_types:
            raise Exception(F"PreservableType with name '{name}' has not been"
                            " registered with Preserve")

        if escaped_state_key in dict_jam:
            obj_state = restore(dict_jam[escaped_state_key])
        else:
            # Object state was a dict and collapsed
            obj_state = restored_dict

        preservable_type = _restore_types[name]
        return preservable_type.restore_method(preservable_type.type, obj_state)

    else:
        # A plain dict
        return restored_dict