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