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
|
# mypy: allow-untyped-defs
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)
|
