Merge branch 'lazily-unpickling-dictionary' into 'master'

Lazily unpickling dictionary

Closes #29

See merge request !89

loopy/codegen/__init__.py

@@ -506,6 +506,12 @@ def generate_code_v2(kernel):
 
     # }}}
 
+    # For faster unpickling in the common case when implemented_domains isn't needed.
+    from loopy.tools import LazilyUnpicklingDictionary
+    codegen_result = codegen_result.copy(
+            implemented_domains=LazilyUnpicklingDictionary(
+                    codegen_result.implemented_domains))
+
     logger.info("%s: generate code: done" % kernel.name)
 
     if CACHING_ENABLED:

loopy/tools.py

@@ -340,6 +340,69 @@ def compute_sccs(graph):
 # }}}
 
 
+# {{{ lazily unpickling dictionary
+
+
+class _PickledObjectWrapper(object):
+    """
+    A class meant to wrap a pickled value (for :class:`LazilyUnpicklingDictionary`).
+    """
+
+    @classmethod
+    def from_object(cls, obj):
+        if isinstance(obj, cls):
+            return obj
+        from pickle import dumps
+        return cls(dumps(obj))
+
+    def __init__(self, objstring):
+        self.objstring = objstring
+
+    def unpickle(self):
+        from pickle import loads
+        return loads(self.objstring)
+
+    def __getstate__(self):
+        return {"objstring": self.objstring}
+
+
+import collections
+
+
+class LazilyUnpicklingDictionary(collections.MutableMapping):
+    """
+    A dictionary-like object which lazily unpickles its values.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self._map = dict(*args, **kwargs)
+
+    def __getitem__(self, key):
+        value = self._map[key]
+        if isinstance(value, _PickledObjectWrapper):
+            value = self._map[key] = value.unpickle()
+        return value
+
+    def __setitem__(self, key, value):
+        self._map[key] = value
+
+    def __delitem__(self, key):
+        del self._map[key]
+
+    def __len__(self):
+        return len(self._map)
+
+    def __iter__(self):
+        return iter(self._map)
+
+    def __getstate__(self):
+        return {"_map": dict(
+            (key, _PickledObjectWrapper.from_object(val))
+            for key, val in six.iteritems(self._map))}
+
+# }}}
+
+
 def is_interned(s):
     return s is None or intern(s) is s
 

loopy/version.py

@@ -32,4 +32,4 @@ except ImportError:
 else:
     _islpy_version = islpy.version.VERSION_TEXT
 
-DATA_MODEL_VERSION = "v55-islpy%s" % _islpy_version
+DATA_MODEL_VERSION = "v58-islpy%s" % _islpy_version

test/test_misc.py

@@ -92,6 +92,58 @@ def test_SetTrie():
         s.add_or_update(set([1, 4]))
 
 
+class PicklableItem(object):
+
+    flags = {"unpickled": False}
+
+    def __getstate__(self):
+        return True
+
+    def __setstate__(self, state):
+        PicklableItem.flags["unpickled"] = True
+
+
+def test_LazilyUnpicklingDictionary():
+    def is_unpickled():
+        return PicklableItem.flags["unpickled"]
+
+    from loopy.tools import LazilyUnpicklingDictionary
+
+    mapping = LazilyUnpicklingDictionary({0: PicklableItem()})
+
+    assert not is_unpickled()
+
+    from pickle import loads, dumps
+
+    pickled_mapping = dumps(mapping)
+
+    # {{{ test lazy loading
+
+    mapping = loads(pickled_mapping)
+    assert not is_unpickled()
+    list(mapping.keys())
+    assert not is_unpickled()
+    assert isinstance(mapping[0], PicklableItem)
+    assert is_unpickled()
+
+    # }}}
+
+    # {{{ test multi round trip
+
+    mapping = loads(dumps(loads(pickled_mapping)))
+    assert isinstance(mapping[0], PicklableItem)
+
+    # }}}
+
+    # {{{ test empty map
+
+    mapping = LazilyUnpicklingDictionary({})
+    mapping = loads(dumps(mapping))
+    assert len(mapping) == 0
+
+    # }}}
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])