Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.10/site-packages/dill/session.py: 15%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1#!/usr/bin/env python
2#
3# Author: Mike McKerns (mmckerns @caltech and @uqfoundation)
4# Author: Leonardo Gama (@leogama)
5# Copyright (c) 2008-2015 California Institute of Technology.
6# Copyright (c) 2016-2025 The Uncertainty Quantification Foundation.
7# License: 3-clause BSD. The full license text is available at:
8# - https://github.com/uqfoundation/dill/blob/master/LICENSE
9"""
10Pickle and restore the intepreter session.
11"""
13__all__ = [
14 'dump_module', 'load_module', 'load_module_asdict',
15 'dump_session', 'load_session' # backward compatibility
16]
18import re
19import os
20import sys
21import warnings
22import pathlib
23import tempfile
25TEMPDIR = pathlib.PurePath(tempfile.gettempdir())
27# Type hints.
28from typing import Optional, Union
30from dill import _dill, Pickler, Unpickler
31from ._dill import (
32 BuiltinMethodType, FunctionType, MethodType, ModuleType, TypeType,
33 _import_module, _is_builtin_module, _is_imported_module, _main_module,
34 _reverse_typemap, __builtin__, UnpicklingError,
35)
37def _module_map():
38 """get map of imported modules"""
39 from collections import defaultdict
40 from types import SimpleNamespace
41 modmap = SimpleNamespace(
42 by_name=defaultdict(list),
43 by_id=defaultdict(list),
44 top_level={},
45 )
46 for modname, module in sys.modules.items():
47 if modname in ('__main__', '__mp_main__') or not isinstance(module, ModuleType):
48 continue
49 if '.' not in modname:
50 modmap.top_level[id(module)] = modname
51 for objname, modobj in module.__dict__.items():
52 modmap.by_name[objname].append((modobj, modname))
53 modmap.by_id[id(modobj)].append((modobj, objname, modname))
54 return modmap
56IMPORTED_AS_TYPES = (ModuleType, TypeType, FunctionType, MethodType, BuiltinMethodType)
57if 'PyCapsuleType' in _reverse_typemap:
58 IMPORTED_AS_TYPES += (_reverse_typemap['PyCapsuleType'],)
59IMPORTED_AS_MODULES = ('ctypes', 'typing', 'subprocess', 'threading',
60 r'concurrent\.futures(\.\w+)?', r'multiprocessing(\.\w+)?')
61IMPORTED_AS_MODULES = tuple(re.compile(x) for x in IMPORTED_AS_MODULES)
63def _lookup_module(modmap, name, obj, main_module):
64 """lookup name or id of obj if module is imported"""
65 for modobj, modname in modmap.by_name[name]:
66 if modobj is obj and sys.modules[modname] is not main_module:
67 return modname, name
68 __module__ = getattr(obj, '__module__', None)
69 if isinstance(obj, IMPORTED_AS_TYPES) or (__module__ is not None
70 and any(regex.fullmatch(__module__) for regex in IMPORTED_AS_MODULES)):
71 for modobj, objname, modname in modmap.by_id[id(obj)]:
72 if sys.modules[modname] is not main_module:
73 return modname, objname
74 return None, None
76def _stash_modules(main_module):
77 modmap = _module_map()
78 newmod = ModuleType(main_module.__name__)
80 imported = []
81 imported_as = []
82 imported_top_level = [] # keep separated for backward compatibility
83 original = {}
84 for name, obj in main_module.__dict__.items():
85 if obj is main_module:
86 original[name] = newmod # self-reference
87 elif obj is main_module.__dict__:
88 original[name] = newmod.__dict__
89 # Avoid incorrectly matching a singleton value in another package (ex.: __doc__).
90 elif any(obj is singleton for singleton in (None, False, True)) \
91 or isinstance(obj, ModuleType) and _is_builtin_module(obj): # always saved by ref
92 original[name] = obj
93 else:
94 source_module, objname = _lookup_module(modmap, name, obj, main_module)
95 if source_module is not None:
96 if objname == name:
97 imported.append((source_module, name))
98 else:
99 imported_as.append((source_module, objname, name))
100 else:
101 try:
102 imported_top_level.append((modmap.top_level[id(obj)], name))
103 except KeyError:
104 original[name] = obj
106 if len(original) < len(main_module.__dict__):
107 newmod.__dict__.update(original)
108 newmod.__dill_imported = imported
109 newmod.__dill_imported_as = imported_as
110 newmod.__dill_imported_top_level = imported_top_level
111 if getattr(newmod, '__loader__', None) is None and _is_imported_module(main_module):
112 # Trick _is_imported_module() to force saving as an imported module.
113 newmod.__loader__ = True # will be discarded by save_module()
114 return newmod
115 else:
116 return main_module
118def _restore_modules(unpickler, main_module):
119 try:
120 for modname, name in main_module.__dict__.pop('__dill_imported'):
121 main_module.__dict__[name] = unpickler.find_class(modname, name)
122 for modname, objname, name in main_module.__dict__.pop('__dill_imported_as'):
123 main_module.__dict__[name] = unpickler.find_class(modname, objname)
124 for modname, name in main_module.__dict__.pop('__dill_imported_top_level'):
125 main_module.__dict__[name] = __import__(modname)
126 except KeyError:
127 pass
129#NOTE: 06/03/15 renamed main_module to main
130def dump_module(
131 filename: Union[str, os.PathLike] = None,
132 module: Optional[Union[ModuleType, str]] = None,
133 refimported: bool = False,
134 **kwds
135) -> None:
136 """Pickle the current state of :py:mod:`__main__` or another module to a file.
138 Save the contents of :py:mod:`__main__` (e.g. from an interactive
139 interpreter session), an imported module, or a module-type object (e.g.
140 built with :py:class:`~types.ModuleType`), to a file. The pickled
141 module can then be restored with the function :py:func:`load_module`.
143 Args:
144 filename: a path-like object or a writable stream. If `None`
145 (the default), write to a named file in a temporary directory.
146 module: a module object or the name of an importable module. If `None`
147 (the default), :py:mod:`__main__` is saved.
148 refimported: if `True`, all objects identified as having been imported
149 into the module's namespace are saved by reference. *Note:* this is
150 similar but independent from ``dill.settings[`byref`]``, as
151 ``refimported`` refers to virtually all imported objects, while
152 ``byref`` only affects select objects.
153 **kwds: extra keyword arguments passed to :py:class:`Pickler()`.
155 Raises:
156 :py:exc:`PicklingError`: if pickling fails.
158 Examples:
160 - Save current interpreter session state:
162 >>> import dill
163 >>> squared = lambda x: x*x
164 >>> dill.dump_module() # save state of __main__ to /tmp/session.pkl
166 - Save the state of an imported/importable module:
168 >>> import dill
169 >>> import pox
170 >>> pox.plus_one = lambda x: x+1
171 >>> dill.dump_module('pox_session.pkl', module=pox)
173 - Save the state of a non-importable, module-type object:
175 >>> import dill
176 >>> from types import ModuleType
177 >>> foo = ModuleType('foo')
178 >>> foo.values = [1,2,3]
179 >>> import math
180 >>> foo.sin = math.sin
181 >>> dill.dump_module('foo_session.pkl', module=foo, refimported=True)
183 - Restore the state of the saved modules:
185 >>> import dill
186 >>> dill.load_module()
187 >>> squared(2)
188 4
189 >>> pox = dill.load_module('pox_session.pkl')
190 >>> pox.plus_one(1)
191 2
192 >>> foo = dill.load_module('foo_session.pkl')
193 >>> [foo.sin(x) for x in foo.values]
194 [0.8414709848078965, 0.9092974268256817, 0.1411200080598672]
196 - Use `refimported` to save imported objects by reference:
198 >>> import dill
199 >>> from html.entities import html5
200 >>> type(html5), len(html5)
201 (dict, 2231)
202 >>> import io
203 >>> buf = io.BytesIO()
204 >>> dill.dump_module(buf) # saves __main__, with html5 saved by value
205 >>> len(buf.getvalue()) # pickle size in bytes
206 71665
207 >>> buf = io.BytesIO()
208 >>> dill.dump_module(buf, refimported=True) # html5 saved by reference
209 >>> len(buf.getvalue())
210 438
212 *Changed in version 0.3.6:* Function ``dump_session()`` was renamed to
213 ``dump_module()``. Parameters ``main`` and ``byref`` were renamed to
214 ``module`` and ``refimported``, respectively.
216 Note:
217 Currently, ``dill.settings['byref']`` and ``dill.settings['recurse']``
218 don't apply to this function.
219 """
220 for old_par, par in [('main', 'module'), ('byref', 'refimported')]:
221 if old_par in kwds:
222 message = "The argument %r has been renamed %r" % (old_par, par)
223 if old_par == 'byref':
224 message += " to distinguish it from dill.settings['byref']"
225 warnings.warn(message + ".", PendingDeprecationWarning)
226 if locals()[par]: # the defaults are None and False
227 raise TypeError("both %r and %r arguments were used" % (par, old_par))
228 refimported = kwds.pop('byref', refimported)
229 module = kwds.pop('main', module)
231 from .settings import settings
232 protocol = settings['protocol']
233 main = module
234 if main is None:
235 main = _main_module
236 elif isinstance(main, str):
237 main = _import_module(main)
238 if not isinstance(main, ModuleType):
239 raise TypeError("%r is not a module" % main)
240 if hasattr(filename, 'write'):
241 file = filename
242 else:
243 if filename is None:
244 filename = str(TEMPDIR/'session.pkl')
245 file = open(filename, 'wb')
246 try:
247 pickler = Pickler(file, protocol, **kwds)
248 pickler._original_main = main
249 if refimported:
250 main = _stash_modules(main)
251 pickler._main = main #FIXME: dill.settings are disabled
252 pickler._byref = False # disable pickling by name reference
253 pickler._recurse = False # disable pickling recursion for globals
254 pickler._session = True # is best indicator of when pickling a session
255 pickler._first_pass = True
256 pickler._main_modified = main is not pickler._original_main
257 pickler.dump(main)
258 finally:
259 if file is not filename: # if newly opened file
260 file.close()
261 return
263# Backward compatibility.
264def dump_session(filename=None, main=None, byref=False, **kwds):
265 warnings.warn("dump_session() has been renamed dump_module()", PendingDeprecationWarning)
266 dump_module(filename, module=main, refimported=byref, **kwds)
267dump_session.__doc__ = dump_module.__doc__
269class _PeekableReader:
270 """lightweight stream wrapper that implements peek()"""
271 def __init__(self, stream):
272 self.stream = stream
273 def read(self, n):
274 return self.stream.read(n)
275 def readline(self):
276 return self.stream.readline()
277 def tell(self):
278 return self.stream.tell()
279 def close(self):
280 return self.stream.close()
281 def peek(self, n):
282 stream = self.stream
283 try:
284 if hasattr(stream, 'flush'): stream.flush()
285 position = stream.tell()
286 stream.seek(position) # assert seek() works before reading
287 chunk = stream.read(n)
288 stream.seek(position)
289 return chunk
290 except (AttributeError, OSError):
291 raise NotImplementedError("stream is not peekable: %r", stream) from None
293def _make_peekable(stream):
294 """return stream as an object with a peek() method"""
295 import io
296 if hasattr(stream, 'peek'):
297 return stream
298 if not (hasattr(stream, 'tell') and hasattr(stream, 'seek')):
299 try:
300 return io.BufferedReader(stream)
301 except Exception:
302 pass
303 return _PeekableReader(stream)
305def _identify_module(file, main=None):
306 """identify the name of the module stored in the given file-type object"""
307 from pickletools import genops
308 UNICODE = {'UNICODE', 'BINUNICODE', 'SHORT_BINUNICODE'}
309 found_import = False
310 try:
311 for opcode, arg, pos in genops(file.peek(256)):
312 if not found_import:
313 if opcode.name in ('GLOBAL', 'SHORT_BINUNICODE') and \
314 arg.endswith('_import_module'):
315 found_import = True
316 else:
317 if opcode.name in UNICODE:
318 return arg
319 else:
320 raise UnpicklingError("reached STOP without finding main module")
321 except (NotImplementedError, ValueError) as error:
322 # ValueError occours when the end of the chunk is reached (without a STOP).
323 if isinstance(error, NotImplementedError) and main is not None:
324 # file is not peekable, but we have main.
325 return None
326 raise UnpicklingError("unable to identify main module") from error
328def load_module(
329 filename: Union[str, os.PathLike] = None,
330 module: Optional[Union[ModuleType, str]] = None,
331 **kwds
332) -> Optional[ModuleType]:
333 """Update the selected module (default is :py:mod:`__main__`) with
334 the state saved at ``filename``.
336 Restore a module to the state saved with :py:func:`dump_module`. The
337 saved module can be :py:mod:`__main__` (e.g. an interpreter session),
338 an imported module, or a module-type object (e.g. created with
339 :py:class:`~types.ModuleType`).
341 When restoring the state of a non-importable module-type object, the
342 current instance of this module may be passed as the argument ``main``.
343 Otherwise, a new instance is created with :py:class:`~types.ModuleType`
344 and returned.
346 Args:
347 filename: a path-like object or a readable stream. If `None`
348 (the default), read from a named file in a temporary directory.
349 module: a module object or the name of an importable module;
350 the module name and kind (i.e. imported or non-imported) must
351 match the name and kind of the module stored at ``filename``.
352 **kwds: extra keyword arguments passed to :py:class:`Unpickler()`.
354 Raises:
355 :py:exc:`UnpicklingError`: if unpickling fails.
356 :py:exc:`ValueError`: if the argument ``main`` and module saved
357 at ``filename`` are incompatible.
359 Returns:
360 A module object, if the saved module is not :py:mod:`__main__` or
361 a module instance wasn't provided with the argument ``main``.
363 Examples:
365 - Save the state of some modules:
367 >>> import dill
368 >>> squared = lambda x: x*x
369 >>> dill.dump_module() # save state of __main__ to /tmp/session.pkl
370 >>>
371 >>> import pox # an imported module
372 >>> pox.plus_one = lambda x: x+1
373 >>> dill.dump_module('pox_session.pkl', module=pox)
374 >>>
375 >>> from types import ModuleType
376 >>> foo = ModuleType('foo') # a module-type object
377 >>> foo.values = [1,2,3]
378 >>> import math
379 >>> foo.sin = math.sin
380 >>> dill.dump_module('foo_session.pkl', module=foo, refimported=True)
382 - Restore the state of the interpreter:
384 >>> import dill
385 >>> dill.load_module() # updates __main__ from /tmp/session.pkl
386 >>> squared(2)
387 4
389 - Load the saved state of an importable module:
391 >>> import dill
392 >>> pox = dill.load_module('pox_session.pkl')
393 >>> pox.plus_one(1)
394 2
395 >>> import sys
396 >>> pox in sys.modules.values()
397 True
399 - Load the saved state of a non-importable module-type object:
401 >>> import dill
402 >>> foo = dill.load_module('foo_session.pkl')
403 >>> [foo.sin(x) for x in foo.values]
404 [0.8414709848078965, 0.9092974268256817, 0.1411200080598672]
405 >>> import math
406 >>> foo.sin is math.sin # foo.sin was saved by reference
407 True
408 >>> import sys
409 >>> foo in sys.modules.values()
410 False
412 - Update the state of a non-importable module-type object:
414 >>> import dill
415 >>> from types import ModuleType
416 >>> foo = ModuleType('foo')
417 >>> foo.values = ['a','b']
418 >>> foo.sin = lambda x: x*x
419 >>> dill.load_module('foo_session.pkl', module=foo)
420 >>> [foo.sin(x) for x in foo.values]
421 [0.8414709848078965, 0.9092974268256817, 0.1411200080598672]
423 *Changed in version 0.3.6:* Function ``load_session()`` was renamed to
424 ``load_module()``. Parameter ``main`` was renamed to ``module``.
426 See also:
427 :py:func:`load_module_asdict` to load the contents of module saved
428 with :py:func:`dump_module` into a dictionary.
429 """
430 if 'main' in kwds:
431 warnings.warn(
432 "The argument 'main' has been renamed 'module'.",
433 PendingDeprecationWarning
434 )
435 if module is not None:
436 raise TypeError("both 'module' and 'main' arguments were used")
437 module = kwds.pop('main')
438 main = module
439 if hasattr(filename, 'read'):
440 file = filename
441 else:
442 if filename is None:
443 filename = str(TEMPDIR/'session.pkl')
444 file = open(filename, 'rb')
445 try:
446 file = _make_peekable(file)
447 #FIXME: dill.settings are disabled
448 unpickler = Unpickler(file, **kwds)
449 unpickler._session = True
451 # Resolve unpickler._main
452 pickle_main = _identify_module(file, main)
453 if main is None and pickle_main is not None:
454 main = pickle_main
455 if isinstance(main, str):
456 if main.startswith('__runtime__.'):
457 # Create runtime module to load the session into.
458 main = ModuleType(main.partition('.')[-1])
459 else:
460 main = _import_module(main)
461 if main is not None:
462 if not isinstance(main, ModuleType):
463 raise TypeError("%r is not a module" % main)
464 unpickler._main = main
465 else:
466 main = unpickler._main
468 # Check against the pickle's main.
469 is_main_imported = _is_imported_module(main)
470 if pickle_main is not None:
471 is_runtime_mod = pickle_main.startswith('__runtime__.')
472 if is_runtime_mod:
473 pickle_main = pickle_main.partition('.')[-1]
474 error_msg = "can't update{} module{} %r with the saved state of{} module{} %r"
475 if is_runtime_mod and is_main_imported:
476 raise ValueError(
477 error_msg.format(" imported", "", "", "-type object")
478 % (main.__name__, pickle_main)
479 )
480 if not is_runtime_mod and not is_main_imported:
481 raise ValueError(
482 error_msg.format("", "-type object", " imported", "")
483 % (pickle_main, main.__name__)
484 )
485 if main.__name__ != pickle_main:
486 raise ValueError(error_msg.format("", "", "", "") % (main.__name__, pickle_main))
488 # This is for find_class() to be able to locate it.
489 if not is_main_imported:
490 runtime_main = '__runtime__.%s' % main.__name__
491 sys.modules[runtime_main] = main
493 loaded = unpickler.load()
494 finally:
495 if not hasattr(filename, 'read'): # if newly opened file
496 file.close()
497 try:
498 del sys.modules[runtime_main]
499 except (KeyError, NameError):
500 pass
501 assert loaded is main
502 _restore_modules(unpickler, main)
503 if main is _main_module or main is module:
504 return None
505 else:
506 return main
508# Backward compatibility.
509def load_session(filename=None, main=None, **kwds):
510 warnings.warn("load_session() has been renamed load_module().", PendingDeprecationWarning)
511 load_module(filename, module=main, **kwds)
512load_session.__doc__ = load_module.__doc__
514def load_module_asdict(
515 filename: Union[str, os.PathLike] = None,
516 update: bool = False,
517 **kwds
518) -> dict:
519 """
520 Load the contents of a saved module into a dictionary.
522 ``load_module_asdict()`` is the near-equivalent of::
524 lambda filename: vars(dill.load_module(filename)).copy()
526 however, does not alter the original module. Also, the path of
527 the loaded module is stored in the ``__session__`` attribute.
529 Args:
530 filename: a path-like object or a readable stream. If `None`
531 (the default), read from a named file in a temporary directory.
532 update: if `True`, initialize the dictionary with the current state
533 of the module prior to loading the state stored at filename.
534 **kwds: extra keyword arguments passed to :py:class:`Unpickler()`
536 Raises:
537 :py:exc:`UnpicklingError`: if unpickling fails
539 Returns:
540 A copy of the restored module's dictionary.
542 Note:
543 If ``update`` is True, the corresponding module may first be imported
544 into the current namespace before the saved state is loaded from
545 filename to the dictionary. Note that any module that is imported into
546 the current namespace as a side-effect of using ``update`` will not be
547 modified by loading the saved module in filename to a dictionary.
549 Example:
550 >>> import dill
551 >>> alist = [1, 2, 3]
552 >>> anum = 42
553 >>> dill.dump_module()
554 >>> anum = 0
555 >>> new_var = 'spam'
556 >>> main = dill.load_module_asdict()
557 >>> main['__name__'], main['__session__']
558 ('__main__', '/tmp/session.pkl')
559 >>> main is globals() # loaded objects don't reference globals
560 False
561 >>> main['alist'] == alist
562 True
563 >>> main['alist'] is alist # was saved by value
564 False
565 >>> main['anum'] == anum # changed after the session was saved
566 False
567 >>> new_var in main # would be True if the option 'update' was set
568 False
569 """
570 if 'module' in kwds:
571 raise TypeError("'module' is an invalid keyword argument for load_module_asdict()")
572 if hasattr(filename, 'read'):
573 file = filename
574 else:
575 if filename is None:
576 filename = str(TEMPDIR/'session.pkl')
577 file = open(filename, 'rb')
578 try:
579 file = _make_peekable(file)
580 main_name = _identify_module(file)
581 old_main = sys.modules.get(main_name)
582 main = ModuleType(main_name)
583 if update:
584 if old_main is None:
585 old_main = _import_module(main_name)
586 main.__dict__.update(old_main.__dict__)
587 else:
588 main.__builtins__ = __builtin__
589 sys.modules[main_name] = main
590 load_module(file, **kwds)
591 finally:
592 if not hasattr(filename, 'read'): # if newly opened file
593 file.close()
594 try:
595 if old_main is None:
596 del sys.modules[main_name]
597 else:
598 sys.modules[main_name] = old_main
599 except NameError: # failed before setting old_main
600 pass
601 main.__session__ = str(filename)
602 return main.__dict__
605# Internal exports for backward compatibility with dill v0.3.5.1
606# Can't be placed in dill._dill because of circular import problems.
607for name in (
608 '_lookup_module', '_module_map', '_restore_modules', '_stash_modules',
609 'dump_session', 'load_session' # backward compatibility functions
610):
611 setattr(_dill, name, globals()[name])
612del name