diff --git a/loopy/codegen/__init__.py b/loopy/codegen/__init__.py index 6f312ec798e13fa4b1d183c27578089857b13e3d..009dadc1a0d6236f092029dbc03ad0c035c7b8f8 100644 --- a/loopy/codegen/__init__.py +++ b/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: diff --git a/loopy/tools.py b/loopy/tools.py index 01d0641fc25c11a092185125604613819a0293ca..56b673b597fc3bf43a6b03f87607ea8d3db0866a 100644 --- a/loopy/tools.py +++ b/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 diff --git a/loopy/version.py b/loopy/version.py index bb2181bfa453a39ab4fb63c3e067de9dea4ab98e..2c15cc2dae6a1ef4b81a43a5b0025f1a0f16fc6a 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 = "v55-islpy%s" % _islpy_version +DATA_MODEL_VERSION = "v58-islpy%s" % _islpy_version diff --git a/test/test_misc.py b/test/test_misc.py index ca4aee5b816eee3eae491c3eecdad296ac04323f..a22e424630255df4225586eeb9f0d62a03d5318f 100644 --- a/test/test_misc.py +++ b/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])