Source code for dantro.data_loaders.pickle
"""Defines a data loader for Python pickles."""
from .._import_tools import LazyLoader
from ..containers import ObjectContainer
from ._registry import add_loader
pkl = LazyLoader("dill")
# -----------------------------------------------------------------------------
[docs]class PickleLoaderMixin:
"""Supplies a load function for pickled python objects.
For unpickling, the ``dill`` package is used.
"""
@add_loader(TargetCls=ObjectContainer, register_aliases=["pkl"])
def _load_pickle(
filepath: str, *, TargetCls: type, **pkl_kwargs
) -> ObjectContainer:
"""Load a pickled object using :py:func:`dill._dill.load`.
Args:
filepath (str): Where the pickle-dumped file is located
TargetCls (type): The class constructor
**pkl_kwargs: Passed on to :py:func:`dill._dill.load`
Returns:
ObjectContainer: The unpickled object, stored in a dantro container
"""
with open(filepath, mode="rb") as f:
obj = pkl.load(f, **pkl_kwargs)
# Populate the target container with the object
return TargetCls(data=obj, attrs=dict(filepath=filepath))