"""
Utils for converting variables into standard python objects (lists, dictionaries, strings, etc.) and back (e.g., for a more predictable LAN transfer).
Provides an extension for pickle for more customized classes (numpy arrays, Dictionary).
"""
import pickle
import collections
import numpy as np
[docs]
class StrDumper:
"""
Class for dumping and loading an object.
Stores procedures for dumping and loading, i.e.,
conversion from complex classes (such as :class:`.Dictionary`) to simple built-in classes (such as :class:`dict` or :class:`str`).
"""
def __init__(self):
self._classes={}
self.add_class(tuple,name="t",loadf=lambda t: tuple(self.load(v) for v in t)) # to distinguish from packed value tuples
_ClassRecord=collections.namedtuple("_ClassRecord",["cls","dump","load","allow_subclass","recursive"])
[docs]
def add_class(self, cls, dumpf=None, loadf=None, name=None, allow_subclass=True, recursive=False):
"""
Add a rule for dumping/loading an object of class `cls`.
Args:
cls
dumpf (callable): Function for dumping an object of the class; ``None`` means identity function.
loadf (callable): Function for loading an object of the class; ``None`` means identity function.
name (str): Name of class, which is stored in the packed data (``cls.__name__`` by default).
allow_subclass (bool): If ``True``, this rule is also used for subclasses of this class.
recursive (bool): If ``True``, the functions are given a second argument, which is a dumping/loading function for their sub-elements.
"""
if name is None:
name=cls.__name__
if name in self._classes:
raise ValueError("class {} is already registered".format(name))
self._classes[name]=self._ClassRecord(cls,dumpf,loadf,allow_subclass,recursive)
def _find_cls(self, obj):
ocls=type(object)
found=None
for n,v in self._classes.items():
cls=v.cls
if v.allow_subclass:
sat=isinstance(obj,cls)
else:
sat=(ocls is cls)
if sat:
if found is None:
found=n,cls
elif issubclass(cls,found[1]): # pylint: disable=unsubscriptable-object
found=n,cls
elif not issubclass(found[1],cls): # pylint: disable=unsubscriptable-object
raise ValueError("both {} and {} satisfy for a base class of {}".format(n,found[0],obj)) # pylint: disable=unsubscriptable-object
return None if found is None else found[0]
def _dump_recursive(self, value):
if isinstance(value,tuple):
return tuple(self.dump(v) for v in value)
elif isinstance(value,list):
return list(self.dump(v) for v in value)
elif isinstance(value,dict):
return dict([(k,self.dump(v)) for k,v in value.items()])
return value
[docs]
def dump(self, obj):
"""Convert an object into a dumped value"""
obj=self._dump_recursive(obj)
reg=self._find_cls(obj)
if reg is None:
return obj
else:
cls=self._classes[reg]
if cls.dump is None:
value=obj
elif not cls.recursive:
value=cls.dump(obj)
else:
value=cls.dump(obj,self.dump)
return (reg,value)
def _load_recursive(self, obj):
if isinstance(obj,tuple):
return tuple(self.load(v) for v in obj)
if isinstance(obj,list):
return list(self.load(v) for v in obj)
if isinstance(obj,dict):
return dict([(k,self.load(v)) for k,v in obj.items()])
return obj
[docs]
def load(self, obj):
"""Convert a dumped value into an object"""
if not isinstance(obj,tuple):
return self._load_recursive(obj)
name,value=obj
if name in self._classes:
cls=self._classes[name]
if cls.load is None:
return value
elif not cls.recursive:
return cls.load(value)
else:
return cls.load(value,self.load)
else:
raise KeyError("class {} is not registered".format(name))
[docs]
def loads(self, s):
"""Convert a pickled string of a damped object into an object"""
value=pickle.loads(s)
return self.load(value)
[docs]
def dumps(self, obj):
"""Dump an object into a pickled string"""
value=self.dump(obj)
return pickle.dumps(value,protocol=-1)
dumper=StrDumper()
"""
Default dumper for converting into standard Python classes and pickling.
Converts :class:`numpy.ndarray` and :class:`.Dictionary` objects
(these conversion routines are defined when corresponding modules are imported).
The converted values include non-printable characters (conversion uses :func:`numpy.load` and :meth:`numpy.ndarray.dump`),
so they can't be saved into text files. However, they're suited for pickling.
"""
### Numpy array ###
dumper.add_class(np.ndarray,np.ndarray.dumps,pickle.loads,"np")
[docs]
def dump(obj):
"""Convert obj into standard Python classes using the default dumper"""
return dumper.dump(obj)
[docs]
def load(s):
"""Convert standard Python class representation `s` into an object using the default dumper"""
return dumper.load(s)
[docs]
def dumps(obj):
"""Convert obj into a pickled string using the default dumper"""
return dumper.dumps(obj)
[docs]
def loads(s):
"""Convert a pickled string into an object using the default dumper"""
return dumper.loads(s)