From c5427abec5605f298bf09b9b044473b5253b70d4 Mon Sep 17 00:00:00 2001
From: Matt Wala <wala1@illinois.edu>
Date: Fri, 17 Feb 2017 15:51:10 -0600
Subject: [PATCH] Implement LazilyUnpicklingDictionary.

---
 loopy/tools.py    | 65 +++++++++++++++++++++++++++++++++++++++++++++++
 test/test_misc.py | 44 ++++++++++++++++++++++++++++++++
 2 files changed, 109 insertions(+)

diff --git a/loopy/tools.py b/loopy/tools.py
index 01d0641fc..36f015a29 100644
--- a/loopy/tools.py
+++ b/loopy/tools.py
@@ -340,6 +340,71 @@ def compute_sccs(graph):
 # }}}
 
 
+# {{{ lazily unpickling dictionary
+
+
+class _PickledObjectWrapper(object):
+    """
+    A class meant to wrap a pickled value (for :class:`LazilyUnpicklingDictionary`).
+    """
+
+    __slots__ = ["objstring"]
+
+    @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)
+
+
+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 dict(
+            (key, _PickledObjectWrapper.from_object(val))
+            for key, val in six.iteritems(self._map))
+
+    def __setstate__(self, state):
+        self._map = state
+
+# }}}
+
+
 def is_interned(s):
     return s is None or intern(s) is s
 
diff --git a/test/test_misc.py b/test/test_misc.py
index ca4aee5b8..139c337e1 100644
--- a/test/test_misc.py
+++ b/test/test_misc.py
@@ -92,6 +92,50 @@ def test_SetTrie():
         s.add_or_update(set([1, 4]))
 
 
+class PicklableItem(object):
+
+    flags = {"unpickled": False}
+
+    def __getstate__(self):
+        return ()
+
+    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)
+
+    # }}}
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])
-- 
GitLab