diff --git a/tests/modelling/test_engine.py b/tests/modelling/test_engine.py index dc9b577f..710a0de4 100644 --- a/tests/modelling/test_engine.py +++ b/tests/modelling/test_engine.py @@ -45,7 +45,7 @@ from zipline.modelling.factor.technical import ( MaxDrawdown, SimpleMovingAverage, ) -from zipline.utils.lazyval import lazyval +from zipline.utils.memoize import lazyval from zipline.utils.test_utils import ( make_rotating_asset_info, make_simple_asset_info, diff --git a/tests/test_doctests.py b/tests/test_doctests.py index 2bc8b35b..d6a5c11a 100644 --- a/tests/test_doctests.py +++ b/tests/test_doctests.py @@ -9,7 +9,7 @@ from zipline.modelling import ( expression, ) from zipline.utils import ( - lazyval, + memoize, test_utils, ) @@ -47,8 +47,8 @@ class DoctestTestCase(TestCase): def test_engine_docs(self): self._check_docs(engine) - def test_lazyval_docs(self): - self._check_docs(lazyval) + def test_memoize_docs(self): + self._check_docs(memoize) def test_test_utils_docs(self): self._check_docs(test_utils) diff --git a/tests/test_memoize.py b/tests/test_memoize.py new file mode 100644 index 00000000..75996e05 --- /dev/null +++ b/tests/test_memoize.py @@ -0,0 +1,34 @@ +""" +Tests for zipline.utils.memoize. +""" +from unittest import TestCase + +from zipline.utils.memoize import remember_last + + +class TestRememberLast(TestCase): + + def test_remember_last(self): + + # Store the count in a list so we can mutate it from inside `func`. + call_count = [0] + + @remember_last + def func(x): + call_count[0] += 1 + return x + + self.assertEqual((func(1), call_count[0]), (1, 1)) + + # Calling again with the same argument should just re-use the old + # value, which means func shouldn't get called again. + self.assertEqual((func(1), call_count[0]), (1, 1)) + self.assertEqual((func(1), call_count[0]), (1, 1)) + + # Calling with a new value should increment the counter. + self.assertEqual((func(2), call_count[0]), (2, 2)) + self.assertEqual((func(2), call_count[0]), (2, 2)) + + # Calling the old value should still increment the counter. + self.assertEqual((func(1), call_count[0]), (1, 3)) + self.assertEqual((func(1), call_count[0]), (1, 3)) diff --git a/zipline/modelling/term.py b/zipline/modelling/term.py index c86d6353..17f880b2 100644 --- a/zipline/modelling/term.py +++ b/zipline/modelling/term.py @@ -15,7 +15,7 @@ from zipline.errors import ( WindowLengthNotPositive, WindowLengthNotSpecified, ) -from zipline.utils.lazyval import lazyval +from zipline.utils.memoize import lazyval NotSpecified = (object(),) diff --git a/zipline/utils/memoize.py b/zipline/utils/memoize.py new file mode 100644 index 00000000..354649bf --- /dev/null +++ b/zipline/utils/memoize.py @@ -0,0 +1,101 @@ +""" +Tools for memoization of function results. +""" +from functools import wraps +from six import iteritems +from weakref import WeakKeyDictionary + + +class lazyval(object): + """ + Decorator that marks that an attribute should not be computed until + needed, and that the value should be memoized. + + Example + ------- + + >>> from zipline.utils.memoize import lazyval + >>> class C(object): + ... def __init__(self): + ... self.count = 0 + ... @lazyval + ... def val(self): + ... self.count += 1 + ... return "val" + ... + >>> c = C() + >>> c.count + 0 + >>> c.val, c.count + ('val', 1) + >>> c.val, c.count + ('val', 1) + """ + def __init__(self, get): + self._get = get + self._cache = WeakKeyDictionary() + + def __get__(self, instance, owner): + if instance is None: + return self + try: + return self._cache[instance] + except KeyError: + self._cache[instance] = val = self._get(instance) + return val + + +def remember_last(f): + """ + Decorator that remembers the last computed value of a function and doesn't + recompute it when called with the same inputs multiple times. + + Parameters + ---------- + f : The function to be memoized. All arguments to f should be hashable. + + Example + ------- + >>> counter = 0 + >>> @remember_last + ... def foo(x): + ... global counter + ... counter += 1 + ... return x, counter + >>> foo(1) + (1, 1) + >>> foo(1) + (1, 1) + >>> foo(0) + (0, 2) + >>> foo(1) + (1, 3) + + Notes + ----- + This decorator is equivalent to `lru_cache(1)` in Python 3, but with less + bells and whistles for handling things like threadsafety. If we ever + decide we need such bells and whistles, we should just make functools32 a + dependency. + """ + # This needs to be a mutable data structure so we can change it from inside + # the function. In pure Python 3, we'd use the nonlocal keyword for this. + _previous = [None, None] + KEY, VALUE = 0, 1 + + _kwd_mark = object() + + @wraps(f) + def memoized_f(*args, **kwds): + # Hashing logic taken from functools32.lru_cache. + key = args + if kwds: + key += _kwd_mark + tuple(sorted(iteritems(kwds))) + + key_hash = hash(key) + if key_hash != _previous[KEY]: + _previous[VALUE] = f(*args, **kwds) + _previous[KEY] = key_hash + return _previous[VALUE] + + return memoized_f