import pickle as pkl
from ..containers import ObjectContainer
from ._tools import add_loader
# -----------------------------------------------------------------------------
[docs]class PickleLoaderMixin:
"""Supplies functionality to load pickled python objects"""
# Define the load function such that it can be set in a flexible way
_PICKLE_LOAD_FUNC = pkl.load
@add_loader(TargetCls=ObjectContainer, omit_self=False)
def _load_pickle(self, filepath: str, *, TargetCls: type, **pkl_kwargs) -> ObjectContainer:
"""Load a pickled object.
This uses the load function defined under the _PICKLE_LOAD_FUNC class
variable, which defaults to the pickle.load function.
Args:
filepath (str): Where the pickle-dumped file is located
TargetCls (type): The class constructor
**pkl_kwargs: Passed on to the load function
Returns:
ObjectContainer: The unpickled file, stored in a dantro container
"""
# Open file in binary mode and unpickle with the given load function
with open(filepath, mode='rb') as f:
obj = self._PICKLE_LOAD_FUNC(f, **pkl_kwargs)
# Populate the target container with the object
return TargetCls(data=obj, attrs=dict(filepath=filepath))
# Also make available under `pkl`
_load_pkl = _load_pickle