Merge pull request #138 from synthicity/memoize-injectable

Injectable Memoization

docs/sim/index.rst

@@ -340,9 +340,21 @@ Use the :py:func:`~urbansim.sim.simulation.add_injectable` function or the
 
 Be default injectable functions are evaluated before injection and the return
 value is passed into other functions. Use ``autocall=False`` to disable this
-behavior and instead inject the wrapped function itself.
-Like tables and columns, injectable functions can have their results
-cached with ``cache=True``.
+behavior and instead inject the function itself.
+Like tables and columns, injectable functions that are automatically evaluated
+can have their results cached with ``cache=True``.
+
+Functions that are not automatically evaluated can also have their results
+cached using the ``memoize=True`` keyword along with ``autocall=False``.
+A memoized injectable will cache results based on the function inputs,
+so this only works if the function inputs are hashable
+(usable as dictionary keys).
+Memoized functions can have their caches cleared manually using their
+``clear_cached`` function attribute.
+The caches of memoized functions are also hooked into the global simulation
+caching system,
+so you can also manage their caches via the ``cache_scope`` keyword argument
+and the :py:func:`~urbansim.sim.simulation.clear_cache` function.
 
 An example of the above injectables in IPython:
 

urbansim/sim/simulation.py

@@ -6,6 +6,7 @@
 import warnings
 from collections import Callable, namedtuple
 from contextlib import contextmanager
+from functools import wraps
 
 import pandas as pd
 import tables
@@ -27,6 +28,7 @@
 _TABLE_CACHE = {}
 _COLUMN_CACHE = {}
 _INJECTABLE_CACHE = {}
+_MEMOIZED = {}
 
 _CS_FOREVER = 'forever'
 _CS_ITER = 'iteration'
@@ -48,6 +50,9 @@ def clear_sim():
     _TABLE_CACHE.clear()
     _COLUMN_CACHE.clear()
     _INJECTABLE_CACHE.clear()
+    for m in _MEMOIZED.values():
+        m.value.clear_cached()
+    _MEMOIZED.clear()
     logger.debug('simulation state cleared')
 
 
@@ -66,12 +71,16 @@ def clear_cache(scope=None):
         _TABLE_CACHE.clear()
         _COLUMN_CACHE.clear()
         _INJECTABLE_CACHE.clear()
+        for m in _MEMOIZED.values():
+            m.value.clear_cached()
         logger.debug('simulation cache cleared')
     else:
         for d in (_TABLE_CACHE, _COLUMN_CACHE, _INJECTABLE_CACHE):
             items = toolz.valfilter(lambda x: x.scope == scope, d)
             for k in items:
                 del d[k]
+        for m in toolz.filter(lambda x: x.scope == scope, _MEMOIZED.values()):
+            m.value.clear_cached()
         logger.debug('cleared cached values with scope {!r}'.format(scope))
 
 
@@ -1022,8 +1031,52 @@ def _columns_for_table(table_name):
             if tname == table_name}
 
 
+def _memoize_function(f, name, cache_scope=_CS_FOREVER):
+    """
+    Wraps a function for memoization and ties it's cache into the
+    simulation cacheing system.
+
+    Parameters
+    ----------
+    f : function
+    name : str
+        Name of injectable.
+    cache_scope : {'step', 'iteration', 'forever'}, optional
+        Scope for which to cache data. Default is to cache forever
+        (or until manually cleared). 'iteration' caches data for each
+        complete iteration of the simulation, 'step' caches data for
+        a single step of the simulation.
+
+    """
+    cache = {}
+
+    @wraps(f)
+    def wrapper(*args, **kwargs):
+        try:
+            cache_key = (
+                args or None, frozenset(kwargs.items()) if kwargs else None)
+            in_cache = cache_key in cache
+        except TypeError:
+            raise TypeError(
+                'function arguments must be hashable for memoization')
+
+        if _CACHING and in_cache:
+            return cache[cache_key]
+        else:
+            result = f(*args, **kwargs)
+            cache[cache_key] = result
+            return result
+
+    wrapper.cache = cache
+    wrapper.clear_cached = lambda: cache.clear()
+    _MEMOIZED[name] = CacheItem(name, wrapper, cache_scope)
+
+    return wrapper
+
+
 def add_injectable(
-        name, value, autocall=True, cache=False, cache_scope=_CS_FOREVER):
+        name, value, autocall=True, cache=False, cache_scope=_CS_FOREVER,
+        memoize=False):
     """
     Add a value that will be injected into other functions.
 
@@ -1049,18 +1102,30 @@ def add_injectable(
         (or until manually cleared). 'iteration' caches data for each
         complete iteration of the simulation, 'step' caches data for
         a single step of the simulation.
+    memoize : bool, optional
+        If autocall is False it is still possible to cache function results
+        by setting this flag to True. Cached values are stored in a dictionary
+        keyed by argument values, so the argument values must be hashable.
+        Memoized functions have their caches cleared according to the same
+        rules as universal caching.
+
+    """
+    if isinstance(value, Callable):
+        if autocall:
+            value = _InjectableFuncWrapper(
+                name, value, cache=cache, cache_scope=cache_scope)
+            # clear any cached data from a previously registered value
+            value.clear_cached()
+        elif not autocall and memoize:
+            value = _memoize_function(value, name, cache_scope=cache_scope)
 
-    """
-    if isinstance(value, Callable) and autocall:
-        value = _InjectableFuncWrapper(
-            name, value, cache=cache, cache_scope=cache_scope)
-        # clear any cached data from a previously registered value
-        value.clear_cached()
     logger.debug('registering injectable {!r}'.format(name))
     _INJECTABLES[name] = value
 
 
-def injectable(name=None, autocall=True, cache=False, cache_scope=_CS_FOREVER):
+def injectable(
+        name=None, autocall=True, cache=False, cache_scope=_CS_FOREVER,
+        memoize=False):
     """
     Decorates functions that will be injected into other functions.
 
@@ -1078,7 +1143,8 @@ def decorator(func):
         else:
             n = func.__name__
         add_injectable(
-            n, func, autocall=autocall, cache=cache, cache_scope=cache_scope)
+            n, func, autocall=autocall, cache=cache, cache_scope=cache_scope,
+            memoize=memoize)
         return func
     return decorator
 

urbansim/sim/tests/test_simulation.py

@@ -598,6 +598,54 @@ def inj():
     assert i()() == 16
 
 
+def test_memoized_injectable():
+    outside = 'x'
+
+    @sim.injectable(autocall=False, memoize=True)
+    def x(s):
+        return outside + s
+
+    assert 'x' in sim._MEMOIZED
+
+    getx = lambda: sim.get_injectable('x')
+
+    assert hasattr(getx(), 'cache')
+    assert hasattr(getx(), 'clear_cached')
+
+    assert getx()('y') == 'xy'
+    outside = 'z'
+    assert getx()('y') == 'xy'
+
+    getx().clear_cached()
+
+    assert getx()('y') == 'zy'
+
+
+def test_memoized_injectable_cache_off():
+    outside = 'x'
+
+    @sim.injectable(autocall=False, memoize=True)
+    def x(s):
+        return outside + s
+
+    getx = lambda: sim.get_injectable('x')('y')
+
+    sim.disable_cache()
+
+    assert getx() == 'xy'
+    outside = 'z'
+    assert getx() == 'zy'
+
+    sim.enable_cache()
+    outside = 'a'
+
+    assert getx() == 'zy'
+
+    sim.disable_cache()
+
+    assert getx() == 'ay'
+
+
 def test_clear_cache_all(df):
     @sim.table(cache=True)
     def table():
@@ -611,18 +659,25 @@ def z(table):
     def x():
         return 'x'
 
+    @sim.injectable(autocall=False, memoize=True)
+    def y(s):
+        return s + 'y'
+
     sim.eval_variable('table.z')
     sim.eval_variable('x')
+    sim.get_injectable('y')('x')
 
     assert sim._TABLE_CACHE.keys() == ['table']
     assert sim._COLUMN_CACHE.keys() == [('table', 'z')]
     assert sim._INJECTABLE_CACHE.keys() == ['x']
+    assert sim._MEMOIZED['y'].value.cache == {(('x',), None): 'xy'}
 
     sim.clear_cache()
 
     assert sim._TABLE_CACHE == {}
     assert sim._COLUMN_CACHE == {}
     assert sim._INJECTABLE_CACHE == {}
+    assert sim._MEMOIZED['y'].value.cache == {}
 
 
 def test_clear_cache_scopes(df):
@@ -638,30 +693,39 @@ def z(table):
     def x():
         return 'x'
 
+    @sim.injectable(autocall=False, memoize=True, cache_scope='iteration')
+    def y(s):
+        return s + 'y'
+
     sim.eval_variable('table.z')
     sim.eval_variable('x')
+    sim.get_injectable('y')('x')
 
     assert sim._TABLE_CACHE.keys() == ['table']
     assert sim._COLUMN_CACHE.keys() == [('table', 'z')]
     assert sim._INJECTABLE_CACHE.keys() == ['x']
+    assert sim._MEMOIZED['y'].value.cache == {(('x',), None): 'xy'}
 
     sim.clear_cache(scope='step')
 
     assert sim._TABLE_CACHE.keys() == ['table']
     assert sim._COLUMN_CACHE.keys() == [('table', 'z')]
     assert sim._INJECTABLE_CACHE == {}
+    assert sim._MEMOIZED['y'].value.cache == {(('x',), None): 'xy'}
 
     sim.clear_cache(scope='iteration')
 
     assert sim._TABLE_CACHE.keys() == ['table']
     assert sim._COLUMN_CACHE == {}
     assert sim._INJECTABLE_CACHE == {}
+    assert sim._MEMOIZED['y'].value.cache == {}
 
     sim.clear_cache(scope='forever')
 
     assert sim._TABLE_CACHE == {}
     assert sim._COLUMN_CACHE == {}
     assert sim._INJECTABLE_CACHE == {}
+    assert sim._MEMOIZED['y'].value.cache == {}
 
 
 def test_cache_scope(df):