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