mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
ed4cda0183
Summary: woooo Test Plan: arc lint --apply-patches --take BLACK --paths-cmd 'hg files -I "caffe2/**/*.py"' Reviewed By: SplitInfinity Differential Revision: D28608934 fbshipit-source-id: 7768fed50a87883a95319376c0a6d73a9492bdcc
27 lines
935 B
Python
27 lines
935 B
Python
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:
|
|
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)
|