_package_unpickler.py 992 B

123456789101112131415161718192021222324252627
  1. # mypy: allow-untyped-defs
  2. import _compat_pickle
  3. import pickle
  4. from .importer import Importer
  5. class PackageUnpickler(pickle._Unpickler): # type: ignore[name-defined]
  6. """Package-aware unpickler.
  7. This behaves the same as a normal unpickler, except it uses `importer` to
  8. find any global names that it encounters while unpickling.
  9. """
  10. def __init__(self, importer: Importer, *args, **kwargs):
  11. super().__init__(*args, **kwargs)
  12. self._importer = importer
  13. def find_class(self, module, name):
  14. # Subclasses may override this.
  15. if self.proto < 3 and self.fix_imports: # type: ignore[attr-defined]
  16. if (module, name) in _compat_pickle.NAME_MAPPING:
  17. module, name = _compat_pickle.NAME_MAPPING[(module, name)]
  18. elif module in _compat_pickle.IMPORT_MAPPING:
  19. module = _compat_pickle.IMPORT_MAPPING[module]
  20. mod = self._importer.import_module(module)
  21. return getattr(mod, name)