diff --git a/loopy/codegen/__init__.py b/loopy/codegen/__init__.py index 009dadc1a0d6236f092029dbc03ad0c035c7b8f8..07bcdc7c6c4a0c23d374a14bc21e4e161b73be03 100644 --- a/loopy/codegen/__init__.py +++ b/loopy/codegen/__init__.py @@ -507,9 +507,9 @@ def generate_code_v2(kernel): # }}} # For faster unpickling in the common case when implemented_domains isn't needed. - from loopy.tools import LazilyUnpicklingDictionary + from loopy.tools import LazilyUnpicklingDict codegen_result = codegen_result.copy( - implemented_domains=LazilyUnpicklingDictionary( + implemented_domains=LazilyUnpicklingDict( codegen_result.implemented_domains)) logger.info("%s: generate code: done" % kernel.name) diff --git a/loopy/tools.py b/loopy/tools.py index 56b673b597fc3bf43a6b03f87607ea8d3db0866a..69a25b375cac5ac519182c71160d4d9b476c4c65 100644 --- a/loopy/tools.py +++ b/loopy/tools.py @@ -23,6 +23,7 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ +import collections import numpy as np from pytools.persistent_dict import KeyBuilder as KeyBuilderBase from loopy.symbolic import WalkMapper as LoopyWalkMapper @@ -340,23 +341,19 @@ def compute_sccs(graph): # }}} -# {{{ lazily unpickling dictionary - +# {{{ pickled container value -class _PickledObjectWrapper(object): - """ - A class meant to wrap a pickled value (for :class:`LazilyUnpicklingDictionary`). +class _PickledObject(object): + """A class meant to wrap a pickled value (for :class:`LazilyUnpicklingDict` and + :class:`LazilyUnpicklingList`). """ - @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 __init__(self, obj): + if isinstance(obj, _PickledObject): + self.objstring = obj.objstring + else: + from pickle import dumps + self.objstring = dumps(obj) def unpickle(self): from pickle import loads @@ -366,12 +363,35 @@ class _PickledObjectWrapper(object): return {"objstring": self.objstring} -import collections +class _PickledObjectWithEqAndPersistentHashKeys(_PickledObject): + """Like :class:`_PickledObject`, with two additional attributes: + * `eq_key` + * `persistent_hash_key` -class LazilyUnpicklingDictionary(collections.MutableMapping): + This allows for comparison and for persistent hashing without unpickling. """ - A dictionary-like object which lazily unpickles its values. + + def __init__(self, obj, eq_key, persistent_hash_key): + _PickledObject.__init__(self, obj) + self.eq_key = eq_key + self.persistent_hash_key = persistent_hash_key + + def update_persistent_hash(self, key_hash, key_builder): + key_builder.rec(key_hash, self.persistent_hash_key) + + def __getstate__(self): + return {"objstring": self.objstring, + "eq_key": self.eq_key, + "persistent_hash_key": self.persistent_hash_key} + +# }}} + + +# {{{ lazily unpickling dictionary + +class LazilyUnpicklingDict(collections.MutableMapping): + """A dictionary-like object which lazily unpickles its values. """ def __init__(self, *args, **kwargs): @@ -379,7 +399,7 @@ class LazilyUnpicklingDictionary(collections.MutableMapping): def __getitem__(self, key): value = self._map[key] - if isinstance(value, _PickledObjectWrapper): + if isinstance(value, _PickledObject): value = self._map[key] = value.unpickle() return value @@ -397,12 +417,105 @@ class LazilyUnpicklingDictionary(collections.MutableMapping): def __getstate__(self): return {"_map": dict( - (key, _PickledObjectWrapper.from_object(val)) + (key, _PickledObject(val)) for key, val in six.iteritems(self._map))} # }}} +# {{{ lazily unpickling list + +class LazilyUnpicklingList(collections.MutableSequence): + """A list which lazily unpickles its values.""" + + def __init__(self, *args, **kwargs): + self._list = list(*args, **kwargs) + + def __getitem__(self, key): + item = self._list[key] + if isinstance(item, _PickledObject): + item = self._list[key] = item.unpickle() + return item + + def __setitem__(self, key, value): + self._list[key] = value + + def __delitem__(self, key): + del self._list[key] + + def __len__(self): + return len(self._list) + + def insert(self, key, value): + self._list.insert(key, value) + + def __getstate__(self): + return {"_list": [_PickledObject(val) for val in self._list]} + + +class LazilyUnpicklingListWithEqAndPersistentHashing(LazilyUnpicklingList): + """A list which lazily unpickles its values, and supports equality comparison + and persistent hashing without unpickling. + + Persistent hashing only works in conjunction with :class:`LoopyKeyBuilder`. + + Equality comparison and persistent hashing are implemented by supplying + functions `eq_key_getter` and `persistent_hash_key_getter` to the + constructor. These functions should return keys that can be used in place of + the original object for the respective purposes of equality comparison and + persistent hashing. + """ + + def __init__(self, *args, **kwargs): + self.eq_key_getter = kwargs.pop("eq_key_getter") + self.persistent_hash_key_getter = kwargs.pop("persistent_hash_key_getter") + LazilyUnpicklingList.__init__(self, *args, **kwargs) + + def update_persistent_hash(self, key_hash, key_builder): + key_builder.update_for_list(key_hash, self._list) + + def _get_eq_key(self, obj): + if isinstance(obj, _PickledObjectWithEqAndPersistentHashKeys): + return obj.eq_key + return self.eq_key_getter(obj) + + def _get_persistent_hash_key(self, obj): + if isinstance(obj, _PickledObjectWithEqAndPersistentHashKeys): + return obj.persistent_hash_key + return self.persistent_hash_key_getter(obj) + + def __eq__(self, other): + if not isinstance(other, (list, LazilyUnpicklingList)): + return NotImplemented + + if isinstance(other, LazilyUnpicklingList): + other = other._list + + if len(self) != len(other): + return False + + for a, b in zip(self._list, other): + if self._get_eq_key(a) != self._get_eq_key(b): + return False + + return True + + def __ne__(self, other): + return not self.__eq__(other) + + def __getstate__(self): + return {"_list": [ + _PickledObjectWithEqAndPersistentHashKeys( + val, + self._get_eq_key(val), + self._get_persistent_hash_key(val)) + for val in self._list], + "eq_key_getter": self.eq_key_getter, + "persistent_hash_key_getter": self.persistent_hash_key_getter} + +# }}} + + def is_interned(s): return s is None or intern(s) is s diff --git a/loopy/version.py b/loopy/version.py index 2e86b974ba224e44cbeb557556a72b82e745525d..f16d689bb823b5e7de09028f29fb91ce668a88e8 100644 --- a/loopy/version.py +++ b/loopy/version.py @@ -32,4 +32,4 @@ except ImportError: else: _islpy_version = islpy.version.VERSION_TEXT -DATA_MODEL_VERSION = "v65-islpy%s" % _islpy_version +DATA_MODEL_VERSION = "v66-islpy%s" % _islpy_version diff --git a/test/test_misc.py b/test/test_misc.py index a22e424630255df4225586eeb9f0d62a03d5318f..0273948b38b28b85e42a600bffb65fbf86dcc554 100644 --- a/test/test_misc.py +++ b/test/test_misc.py @@ -92,26 +92,36 @@ def test_SetTrie(): s.add_or_update(set([1, 4])) -class PicklableItem(object): +class PickleDetector(object): + """Contains a class attribute which flags if any instance was unpickled. + """ - flags = {"unpickled": False} + @classmethod + def reset(cls): + cls.instance_unpickled = False def __getstate__(self): - return True + return {"state": self.state} def __setstate__(self, state): - PicklableItem.flags["unpickled"] = True + self.__class__.instance_unpickled = True + self.state = state["state"] -def test_LazilyUnpicklingDictionary(): - def is_unpickled(): - return PicklableItem.flags["unpickled"] +class PickleDetectorForLazilyUnpicklingDict(PickleDetector): + instance_unpickled = False - from loopy.tools import LazilyUnpicklingDictionary + def __init__(self): + self.state = None - mapping = LazilyUnpicklingDictionary({0: PicklableItem()}) - assert not is_unpickled() +def test_LazilyUnpicklingDict(): + from loopy.tools import LazilyUnpicklingDict + + cls = PickleDetectorForLazilyUnpicklingDict + mapping = LazilyUnpicklingDict({0: cls()}) + + assert not cls.instance_unpickled from pickle import loads, dumps @@ -120,30 +130,160 @@ def test_LazilyUnpicklingDictionary(): # {{{ test lazy loading mapping = loads(pickled_mapping) - assert not is_unpickled() + assert not cls.instance_unpickled list(mapping.keys()) - assert not is_unpickled() - assert isinstance(mapping[0], PicklableItem) - assert is_unpickled() + assert not cls.instance_unpickled + assert isinstance(mapping[0], cls) + assert cls.instance_unpickled + + # }}} + + # {{{ conversion + + cls.reset() + mapping = loads(pickled_mapping) + dict(mapping) + assert cls.instance_unpickled # }}} # {{{ test multi round trip mapping = loads(dumps(loads(pickled_mapping))) - assert isinstance(mapping[0], PicklableItem) + assert isinstance(mapping[0], cls) # }}} # {{{ test empty map - mapping = LazilyUnpicklingDictionary({}) + mapping = LazilyUnpicklingDict({}) mapping = loads(dumps(mapping)) assert len(mapping) == 0 # }}} +class PickleDetectorForLazilyUnpicklingList(PickleDetector): + instance_unpickled = False + + def __init__(self): + self.state = None + + +def test_LazilyUnpicklingList(): + from loopy.tools import LazilyUnpicklingList + + cls = PickleDetectorForLazilyUnpicklingList + lst = LazilyUnpicklingList([cls()]) + assert not cls.instance_unpickled + + from pickle import loads, dumps + pickled_lst = dumps(lst) + + # {{{ test lazy loading + + lst = loads(pickled_lst) + assert not cls.instance_unpickled + assert isinstance(lst[0], cls) + assert cls.instance_unpickled + + # }}} + + # {{{ conversion + + cls.reset() + lst = loads(pickled_lst) + list(lst) + assert cls.instance_unpickled + + # }}} + + # {{{ test multi round trip + + lst = loads(dumps(loads(dumps(lst)))) + assert isinstance(lst[0], cls) + + # }}} + + # {{{ test empty list + + lst = LazilyUnpicklingList([]) + lst = loads(dumps(lst)) + assert len(lst) == 0 + + # }}} + + +class PickleDetectorForLazilyUnpicklingListWithEqAndPersistentHashing( + PickleDetector): + instance_unpickled = False + + def __init__(self, comparison_key): + self.state = comparison_key + + def __repr__(self): + return repr(self.state) + + def update_persistent_hash(self, key_hash, key_builder): + key_builder.rec(key_hash, repr(self)) + + +def test_LazilyUnpicklingListWithEqAndPersistentHashing(): + from loopy.tools import LazilyUnpicklingListWithEqAndPersistentHashing + + cls = PickleDetectorForLazilyUnpicklingListWithEqAndPersistentHashing + from pickle import loads, dumps + + # {{{ test comparison of a pair of lazy lists + + lst0 = LazilyUnpicklingListWithEqAndPersistentHashing( + [cls(0), cls(1)], + eq_key_getter=repr, + persistent_hash_key_getter=repr) + lst1 = LazilyUnpicklingListWithEqAndPersistentHashing( + [cls(0), cls(1)], + eq_key_getter=repr, + persistent_hash_key_getter=repr) + + assert not cls.instance_unpickled + + assert lst0 == lst1 + assert not cls.instance_unpickled + + lst0 = loads(dumps(lst0)) + lst1 = loads(dumps(lst1)) + + assert lst0 == lst1 + assert not cls.instance_unpickled + + lst0.append(cls(3)) + lst1.append(cls(2)) + + assert lst0 != lst1 + + # }}} + + # {{{ comparison with plain lists + + lst = [cls(0), cls(1), cls(3)] + + assert lst == lst0 + assert lst0 == lst + assert not cls.instance_unpickled + + # }}} + + # {{{ persistent hashing + + from loopy.tools import LoopyKeyBuilder + kb = LoopyKeyBuilder() + + assert kb(lst0) == kb(lst) + assert not cls.instance_unpickled + + # }}} + + if __name__ == "__main__": if len(sys.argv) > 1: exec(sys.argv[1])