Change memorise to diff, add turn on function and extra test

dill/__init__.py

@@ -24,8 +24,8 @@
 """ + __license__
 
 from .dill import dump, dumps, load, loads, dump_session, load_session, \
-    Pickler, Unpickler, register, copy, pickle, pickles, HIGHEST_PROTOCOL, \
-    DEFAULT_PROTOCOL, PicklingError, UnpicklingError, \
+    Pickler, Unpickler, register, copy, pickle, pickles, use_diff, \
+    HIGHEST_PROTOCOL, DEFAULT_PROTOCOL, PicklingError, UnpicklingError, \
     FMODE_NEWHANDLE, FMODE_PRESERVEDATA, FMODE_PICKLECONTENTS
 from . import source, temp, detect
 

dill/dill.py

@@ -16,7 +16,7 @@
 """
 __all__ = ['dump','dumps','load','loads','dump_session','load_session',
            'Pickler','Unpickler','register','copy','pickle','pickles',
-           'HIGHEST_PROTOCOL','DEFAULT_PROTOCOL',
+           'use_diff', 'HIGHEST_PROTOCOL','DEFAULT_PROTOCOL',
            'PicklingError','UnpicklingError','FMODE_NEWHANDLE',
            'FMODE_PRESERVEDATA','FMODE_PICKLECONTENTS']
 
@@ -31,10 +31,8 @@ def _trace(boolean):
 
 import os
 import sys
-try:
-    from . import memorise
-except ImportError:
-    import memorise
+diff = None
+_use_diff = False
 PY3 = (hex(sys.hexversion) >= '0x30000f0')
 if PY3: #XXX: get types from dill.objtypes ?
     import builtins as __builtin__
@@ -257,7 +255,7 @@ class Pickler(StockPickler):
     def __init__(self, *args, **kwargs):
         StockPickler.__init__(self, *args, **kwargs)
         self._main_module = _main_module
-        self._memorise_cashe = {}
+        self._diff_cashe = {}
 
 class Unpickler(StockUnpickler):
     """python's Unpickler extended to interpreter sessions and more types"""
@@ -300,6 +298,16 @@ def _revert_extension():
             if type in pickle_dispatch_copy:
                 StockPickler.dispatch[type] = pickle_dispatch_copy[type]
 
+def use_diff(on=True):
+    global _use_diff, diff
+    _use_diff = on
+    if _use_diff and diff is None:
+        try:
+            from . import diff as d
+        except:
+            import diff as d
+        diff = d
+
 def _create_typemap():
     import types
     if PY3:
@@ -890,20 +898,38 @@ def save_weakproxy(pickler, obj):
 
 @register(ModuleType)
 def save_module(pickler, obj):
-    if obj.__name__ != "dill":
-        try:
-            changed = memorise.whats_changed(obj,
-                                             seen=pickler._memorise_cashe)[0]
-        except RuntimeError:  # not memorised module, probably part of dill
-            pass
-        else:
+    if _use_diff:
+        if obj.__name__ != "dill":
+            try:
+                changed = diff.whats_changed(obj, seen=pickler._diff_cashe)[0]
+            except RuntimeError:  # not memorised module, probably part of dill
+                pass
+            else:
+                log.info("M1: %s with diff" % obj)
+                log.info("Diff: %s", changed.keys())
+                pickler.save_reduce(_import_module, (obj.__name__,), obj=obj,
+                                    state=changed)
+                return
+
+        log.info("M2: %s" % obj)
+        pickler.save_reduce(_import_module, (obj.__name__,), obj=obj)
+    else:
+        # if a module file name starts with prefx, it should be a builtin
+        # module, so should be pickled as a reference
+        prefix = getattr(sys, "base_prefix", sys.prefix)
+        std_mod = getattr(obj, "__file__", prefix).startswith(prefix)
+        if obj.__name__ not in ("builtins", "dill") \
+           and not std_mod or is_dill(pickler) and obj is pickler._main_module:
             log.info("M1: %s" % obj)
+            _main_dict = obj.__dict__.copy() #XXX: better no copy? option to copy?
+            [_main_dict.pop(item, None) for item in singletontypes
+                + ["__builtins__", "__loader__"]]
             pickler.save_reduce(_import_module, (obj.__name__,), obj=obj,
-                                state=changed)
-            return
-
-    log.info("M2: %s" % obj)
-    pickler.save_reduce(_import_module, (obj.__name__,), obj=obj)
+                                state=_main_dict)
+        else:
+            log.info("M2: %s" % obj)
+            pickler.save_reduce(_import_module, (obj.__name__,), obj=obj)
+        return
     return
 
 @register(TypeType)

tests/test_memo.py → tests/test_diff.py

@@ -5,7 +5,7 @@
 # License: 3-clause BSD.  The full license text is available at:
 #  - http://trac.mystic.cacr.caltech.edu/project/pathos/browser/dill/LICENSE
 
-from dill import memorise as m
+from dill import diff
 
 
 class A:
@@ -16,47 +16,47 @@ class A:
 c = A()
 a.a = b
 b.a = c
-m.memorise(a)
-assert not m.has_changed(a)
+diff.memorise(a)
+assert not diff.has_changed(a)
 c.a = 1
-assert m.has_changed(a)
-m.memorise(c, force=True)
-assert not m.has_changed(a)
+assert diff.has_changed(a)
+diff.memorise(c, force=True)
+assert not diff.has_changed(a)
 c.a = 2
-assert m.has_changed(a)
-changed = m.whats_changed(a)
+assert diff.has_changed(a)
+changed = diff.whats_changed(a)
 assert list(changed[0].keys()) == ["a"]
 assert not changed[1]
 
 a2 = []
 b2 = [a2]
 c2 = [b2]
-m.memorise(c2)
-assert not m.has_changed(c2)
+diff.memorise(c2)
+assert not diff.has_changed(c2)
 a2.append(1)
-assert m.has_changed(c2)
-changed = m.whats_changed(c2)
+assert diff.has_changed(c2)
+changed = diff.whats_changed(c2)
 assert changed[0] == {}
 assert changed[1]
 
 a3 = {}
 b3 = {1: a3}
 c3 = {1: b3}
-m.memorise(c3)
-assert not m.has_changed(c3)
+diff.memorise(c3)
+assert not diff.has_changed(c3)
 a3[1] = 1
-assert m.has_changed(c3)
-changed = m.whats_changed(c3)
+assert diff.has_changed(c3)
+changed = diff.whats_changed(c3)
 assert changed[0] == {}
 assert changed[1]
 
 import abc
 # make sure that the "_abc_invaldation_counter" does not cause the test to fail
-m.memorise(abc.ABCMeta, force=True)
-assert not m.has_changed(abc)
+diff.memorise(abc.ABCMeta, force=True)
+assert not diff.has_changed(abc)
 abc.ABCMeta.zzz = 1
-assert m.has_changed(abc)
-changed = m.whats_changed(abc)
+assert diff.has_changed(abc)
+changed = diff.whats_changed(abc)
 assert list(changed[0].keys()) == ["ABCMeta"]
 assert not changed[1]
 
@@ -66,14 +66,14 @@ class A:
 c = A()
 a.a = b
 b.a = c
-m.memorise(a)
-assert not m.has_changed(a)
+diff.memorise(a)
+assert not diff.has_changed(a)
 c.a = 1
-assert m.has_changed(a)
-m.memorise(c, force=True)
-assert not m.has_changed(a)
+assert diff.has_changed(a)
+diff.memorise(c, force=True)
+assert not diff.has_changed(a)
 del c.a
-assert m.has_changed(a)
-changed = m.whats_changed(a)
+assert diff.has_changed(a)
+changed = diff.whats_changed(a)
 assert list(changed[0].keys()) == ["a"]
 assert not changed[1]

tests/test_module.py

@@ -8,9 +8,30 @@
 import sys
 import dill
 import test_mixins as module
+import imp
 
 cached = (module.__cached__ if hasattr(module, "__cached__")
-          else module.__file__ + "c")
+          else module.__file__.split(".", 1)[0] + ".pyc")
+
+module.a = 1234
+
+pik_mod = dill.dumps(module)
+
+module.a = 0
+
+# remove module
+del sys.modules[module.__name__]
+del module
+
+module = dill.loads(pik_mod)
+assert hasattr(module, "a") and module.a == 1234
+assert module.double_add(1, 2, 3) == 2 * module.fx
+
+# Restart, and test use_diff
+
+imp.reload(module)
+
+dill.use_diff()
 
 module.a = 1234
 
@@ -29,5 +50,6 @@
 # clean up
 import os
 os.remove(cached)
-if os.path.exists("__pycache__") and not os.listdir("__pycache__"):
-    os.removedirs("__pycache__")
+pycache = os.path.join(os.path.dirname(module.__file__), "__pycache__")
+if os.path.exists(pycache) and not os.listdir(pycache):
+    os.removedirs(pycache)