_package_unpickler.py 965 B

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