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