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