Skip to content

Commit

Permalink
new function is_module_pickle()
Browse files Browse the repository at this point in the history
  • Loading branch information
leogama committed Jul 16, 2022
1 parent da4cc07 commit 1732e3d
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 9 deletions.
8 changes: 4 additions & 4 deletions dill/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,10 +25,10 @@

from ._dill import (
dump, dumps, load, loads, dump_module, load_module, load_module_asdict,
dump_session, load_session, Pickler, Unpickler, register, copy, pickle,
pickles, check, HIGHEST_PROTOCOL, DEFAULT_PROTOCOL, PicklingError,
UnpicklingError, HANDLE_FMODE, CONTENTS_FMODE, FILE_FMODE, PickleError,
PickleWarning, PicklingWarning, UnpicklingWarning,
dump_session, load_session, is_module_pickle, Pickler, Unpickler, register,
copy, pickle, pickles, check, HIGHEST_PROTOCOL, DEFAULT_PROTOCOL,
PicklingError, UnpicklingError, HANDLE_FMODE, CONTENTS_FMODE, FILE_FMODE,
PickleError, PickleWarning, PicklingWarning, UnpicklingWarning,
)
from . import source, temp, detect

Expand Down
32 changes: 27 additions & 5 deletions dill/_dill.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,11 @@
"""
__all__ = [
'dump', 'dumps', 'load', 'loads', 'dump_module', 'load_module',
'load_module_asdict', 'dump_session', 'load_session', 'Pickler', 'Unpickler',
'register', 'copy', 'pickle', 'pickles', 'check', 'HIGHEST_PROTOCOL',
'DEFAULT_PROTOCOL', 'PicklingError', 'UnpicklingError', 'HANDLE_FMODE',
'CONTENTS_FMODE', 'FILE_FMODE', 'PickleError', 'PickleWarning',
'PicklingWarning', 'UnpicklingWarning'
'load_module_asdict', 'dump_session', 'load_session', 'is_module_pickle',
'Pickler', 'Unpickler', 'register', 'copy', 'pickle', 'pickles', 'check',
'HIGHEST_PROTOCOL', 'DEFAULT_PROTOCOL', 'PicklingError', 'UnpicklingError',
'HANDLE_FMODE', 'CONTENTS_FMODE', 'FILE_FMODE', 'PickleError',
'PickleWarning', 'PicklingWarning', 'UnpicklingWarning',
]

__module__ = 'dill'
Expand Down Expand Up @@ -606,6 +606,28 @@ def _identify_module(file, main=None):
return None
raise UnpicklingError("unable to identify main module") from error

def is_module_pickle(filename, importable: bool = True) -> bool:
"""Check if a file is a module state pickle file.
Parameters:
filename: a path-like object or a readable stream.
importable: expected kind of the file's saved module. Use `True` for
importable modules (the default) or `False` for module-type objects.
Returns:
`True` if the pickle file at ``filename`` was generated with
:py:func:`dump_module` **AND** the module whose state is saved in it is
of the kind specified by the ``importable`` argument. `False` otherwise.
"""
with _open(filename, 'rb', peekable=True) as file:
try:
pickle_main = _identify_module(file)
except UnpicklingError:
return False
else:
is_runtime_mod = pickle_main.startswith('__runtime__.')
return importable ^ is_runtime_mod

def load_module(
filename = str(TEMPDIR/'session.pkl'),
module: Union[ModuleType, str] = None,
Expand Down

0 comments on commit 1732e3d

Please sign in to comment.