def test_intern_error(self):
"""Failure to handle all dependencies should lead to an error."""
import package_a.subpackage
obj = package_a.subpackage.PackageASubpackageObject()
obj2 = package_a.PackageAObject(obj)
buffer = BytesIO()
with self.assertRaises(PackagingError) as e:
with PackageExporter(buffer) as he:
he.save_pickle("obj", "obj.pkl", obj2)
self.assertEqual(
str(e.exception),
dedent("""
* Module did not match against any action pattern. Extern, mock, or intern it.
package_a
package_a.subpackage
"""),
)
# Interning all dependencies should work
with PackageExporter(buffer) as he:
he.intern(["package_a", "package_a.subpackage"])
he.save_pickle("obj", "obj.pkl", obj2)
def test_pickle_mocked(self):
import package_a.subpackage
obj = package_a.subpackage.PackageASubpackageObject()
obj2 = package_a.PackageAObject(obj)
filename = self.temp()
with PackageExporter(filename, verbose=False) as he:
he.mock(include='package_a.subpackage')
he.save_pickle('obj', 'obj.pkl', obj2)
hi = PackageImporter(filename)
with self.assertRaises(NotImplementedError):
hi.load_pickle('obj', 'obj.pkl')
def test_unique_module_names(self):
import package_a.subpackage
obj = package_a.subpackage.PackageASubpackageObject()
obj2 = package_a.PackageAObject(obj)
f1 = self.temp()
with PackageExporter(f1, verbose=False) as pe:
pe.save_pickle("obj", "obj.pkl", obj2)
importer1 = PackageImporter(f1)
loaded1 = importer1.load_pickle("obj", "obj.pkl")
importer2 = PackageImporter(f1)
loaded2 = importer2.load_pickle("obj", "obj.pkl")
# Modules from loaded packages should not shadow the names of modules.
# See mangling.md for more info.
self.assertNotEqual(type(obj2).__module__, type(loaded1).__module__)
self.assertNotEqual(type(loaded1).__module__, type(loaded2).__module__)
def test_exporting_mismatched_code(self):
"""
If an object with the same qualified name is loaded from different
packages, the user should get an error if they try to re-save the
object with the wrong package's source code.
"""
import package_a.subpackage
obj = package_a.subpackage.PackageASubpackageObject()
obj2 = package_a.PackageAObject(obj)
f1 = self.temp()
with PackageExporter(f1, verbose=False) as pe:
pe.intern("**")
pe.save_pickle("obj", "obj.pkl", obj2)
importer1 = PackageImporter(f1)
loaded1 = importer1.load_pickle("obj", "obj.pkl")
importer2 = PackageImporter(f1)
loaded2 = importer2.load_pickle("obj", "obj.pkl")
f2 = self.temp()
def make_exporter():
pe = PackageExporter(f2,
verbose=False,
importer=[importer1, sys_importer])
# Ensure that the importer finds the 'PackageAObject' defined in 'importer1' first.
return pe
# This should fail. The 'PackageAObject' type defined from 'importer1'
# is not necessarily the same 'obj2's version of 'PackageAObject'.
pe = make_exporter()
with self.assertRaises(pickle.PicklingError):
pe.save_pickle("obj", "obj.pkl", obj2)
# This should also fail. The 'PackageAObject' type defined from 'importer1'
# is not necessarily the same as the one defined from 'importer2'
pe = make_exporter()
with self.assertRaises(pickle.PicklingError):
pe.save_pickle("obj", "obj.pkl", loaded2)
# This should succeed. The 'PackageAObject' type defined from
# 'importer1' is a match for the one used by loaded1.
pe = make_exporter()
pe.save_pickle("obj", "obj.pkl", loaded1)
def test_save_imported_module_fails(self):
"""
Directly saving/requiring an PackageImported module should raise a specific error message.
"""
import package_a.subpackage
obj = package_a.subpackage.PackageASubpackageObject()
obj2 = package_a.PackageAObject(obj)
f1 = self.temp()
with PackageExporter(f1, verbose=False) as pe:
pe.intern("**")
pe.save_pickle("obj", "obj.pkl", obj)
importer1 = PackageImporter(f1)
loaded1 = importer1.load_pickle("obj", "obj.pkl")
f2 = self.temp()
pe = PackageExporter(f2, verbose=False, importer=(importer1, sys_importer))
with self.assertRaisesRegex(ModuleNotFoundError, "torch.package"):
pe.save_module(loaded1.__module__)
def test_pickle(self):
import package_a.subpackage
obj = package_a.subpackage.PackageASubpackageObject()
obj2 = package_a.PackageAObject(obj)
filename = self.temp()
with PackageExporter(filename, verbose=False) as he:
he.save_pickle('obj', 'obj.pkl', obj2)
hi = PackageImporter(filename)
# check we got dependencies
sp = hi.import_module('package_a.subpackage')
# check we didn't get other stuff
with self.assertRaises(ImportError):
hi.import_module('module_a')
obj_loaded = hi.load_pickle('obj', 'obj.pkl')
self.assertIsNot(obj2, obj_loaded)
self.assertIsInstance(obj_loaded.obj, sp.PackageASubpackageObject)
self.assertIsNot(package_a.subpackage.PackageASubpackageObject, sp.PackageASubpackageObject)
def test_intern_error(self):
"""Failure to handle all dependencies should lead to an error."""
import package_a.subpackage
obj = package_a.subpackage.PackageASubpackageObject()
obj2 = package_a.PackageAObject(obj)
buffer = BytesIO()
try:
with PackageExporter(buffer, verbose=False) as he:
he.save_pickle("obj", "obj.pkl", obj2)
except PackagingError as e:
self.assertEqual(e.unhandled,
set(["package_a", "package_a.subpackage"]))
else:
self.fail("PackagingError should have been raised")
# Interning all dependencies should work
with PackageExporter(buffer, verbose=False) as he:
he.intern(["package_a", "package_a.subpackage"])
he.save_pickle("obj", "obj.pkl", obj2)
def test_save_imported_module(self):
"""Saving a module that came from another PackageImporter should work."""
import package_a.subpackage
obj = package_a.subpackage.PackageASubpackageObject()
obj2 = package_a.PackageAObject(obj)
buffer = BytesIO()
with PackageExporter(buffer) as exporter:
exporter.intern("**")
exporter.save_pickle("model", "model.pkl", obj2)
buffer.seek(0)
importer = PackageImporter(buffer)
imported_obj2 = importer.load_pickle("model", "model.pkl")
imported_obj2_module = imported_obj2.__class__.__module__
# Should export without error.
buffer2 = BytesIO()
with PackageExporter(buffer2, importer=(importer, sys_importer)) as exporter:
exporter.intern("**")
exporter.save_module(imported_obj2_module)
def test_pickle(self):
import package_a.subpackage
obj = package_a.subpackage.PackageASubpackageObject()
obj2 = package_a.PackageAObject(obj)
filename = self.temp()
with PackageExporter(filename) as he:
he.intern("**")
he.save_pickle("obj", "obj.pkl", obj2)
hi = PackageImporter(filename)
# check we got dependencies
sp = hi.import_module("package_a.subpackage")
# check we didn't get other stuff
with self.assertRaises(ImportError):
hi.import_module("module_a")
obj_loaded = hi.load_pickle("obj", "obj.pkl")
self.assertIsNot(obj2, obj_loaded)
self.assertIsInstance(obj_loaded.obj, sp.PackageASubpackageObject)
self.assertIsNot(package_a.subpackage.PackageASubpackageObject,
sp.PackageASubpackageObject)