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

266 statements  

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""" 

12 

13__all__ = [ 

14 'dump_module', 'load_module', 'load_module_asdict', 

15 'dump_session', 'load_session' # backward compatibility 

16] 

17 

18import re 

19import os 

20import sys 

21import warnings 

22 

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) 

29 

30# Type hints. 

31from typing import Optional, Union 

32 

33import pathlib 

34import tempfile 

35 

36TEMPDIR = pathlib.PurePath(tempfile.gettempdir()) 

37 

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 

56 

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) 

63 

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 

76 

77def _stash_modules(main_module): 

78 modmap = _module_map() 

79 newmod = ModuleType(main_module.__name__) 

80 

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 

106 

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 

118 

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 

129 

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. 

138 

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`. 

143 

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()`. 

155 

156 Raises: 

157 :py:exc:`PicklingError`: if pickling fails. 

158 

159 Examples: 

160 

161 - Save current interpreter session state: 

162 

163 >>> import dill 

164 >>> squared = lambda x: x*x 

165 >>> dill.dump_module() # save state of __main__ to /tmp/session.pkl 

166 

167 - Save the state of an imported/importable module: 

168 

169 >>> import dill 

170 >>> import pox 

171 >>> pox.plus_one = lambda x: x+1 

172 >>> dill.dump_module('pox_session.pkl', module=pox) 

173 

174 - Save the state of a non-importable, module-type object: 

175 

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) 

183 

184 - Restore the state of the saved modules: 

185 

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] 

196 

197 - Use `refimported` to save imported objects by reference: 

198 

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 

212 

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. 

216 

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) 

231 

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 

263 

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__ 

269 

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 

293 

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) 

305 

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 

328 

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``. 

336 

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`). 

341 

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. 

346 

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()`. 

354 

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. 

359 

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``. 

363 

364 Examples: 

365 

366 - Save the state of some modules: 

367 

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) 

382 

383 - Restore the state of the interpreter: 

384 

385 >>> import dill 

386 >>> dill.load_module() # updates __main__ from /tmp/session.pkl 

387 >>> squared(2) 

388 4 

389 

390 - Load the saved state of an importable module: 

391 

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 

399 

400 - Load the saved state of a non-importable module-type object: 

401 

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 

412 

413 - Update the state of a non-importable module-type object: 

414 

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] 

423 

424 *Changed in version 0.3.6:* Function ``load_session()`` was renamed to 

425 ``load_module()``. Parameter ``main`` was renamed to ``module``. 

426 

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 

451 

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 

468 

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)) 

488 

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 

493 

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 

508 

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__ 

514 

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. 

522 

523 ``load_module_asdict()`` is the near-equivalent of:: 

524 

525 lambda filename: vars(dill.load_module(filename)).copy() 

526 

527 however, does not alter the original module. Also, the path of 

528 the loaded module is stored in the ``__session__`` attribute. 

529 

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()` 

536 

537 Raises: 

538 :py:exc:`UnpicklingError`: if unpickling fails 

539 

540 Returns: 

541 A copy of the restored module's dictionary. 

542 

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. 

549 

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__ 

604 

605 

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