""" 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