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