pytorch

1
import _compat_pickle
2
import pickle
3

4
from .importer import Importer
5

6

7
class PackageUnpickler(pickle._Unpickler):  # type: ignore[name-defined]
8
    """Package-aware unpickler.
9

10
    This behaves the same as a normal unpickler, except it uses `importer` to
11
    find any global names that it encounters while unpickling.
12
    """
13

14
    def __init__(self, importer: Importer, *args, **kwargs):
15
        super().__init__(*args, **kwargs)
16
        self._importer = importer
17

18
    def find_class(self, module, name):
19
        # Subclasses may override this.
20
        if self.proto < 3 and self.fix_imports:  # type: ignore[attr-defined]
21
            if (module, name) in _compat_pickle.NAME_MAPPING:
22
                module, name = _compat_pickle.NAME_MAPPING[(module, name)]
23
            elif module in _compat_pickle.IMPORT_MAPPING:
24
                module = _compat_pickle.IMPORT_MAPPING[module]
25
        mod = self._importer.import_module(module)
26
        return getattr(mod, name)
27