From 958d455a7a52e51f241e3e8286f5dbd6228143f1 Mon Sep 17 00:00:00 2001 From: Joe Jevnik Date: Sat, 4 Jun 2016 17:57:06 -0400 Subject: [PATCH] ENH: Support default params for terms --- tests/pipeline/test_term.py | 55 +++++++++++++++++++++++++++++++--- zipline/assets/symbol_rules.py | 13 ++++++++ zipline/pipeline/term.py | 15 +++++++--- zipline/testing/predicates.py | 43 ++++++++++++++++++++++++++ 4 files changed, 118 insertions(+), 8 deletions(-) create mode 100644 zipline/assets/symbol_rules.py diff --git a/tests/pipeline/test_term.py b/tests/pipeline/test_term.py index eb97964e..0cda474d 100644 --- a/tests/pipeline/test_term.py +++ b/tests/pipeline/test_term.py @@ -5,6 +5,8 @@ from collections import Counter from itertools import product from unittest import TestCase +from toolz import assoc + from zipline.assets import Asset from zipline.errors import ( DTypeNotSpecified, @@ -31,7 +33,12 @@ from zipline.pipeline.factors import RecarrayField from zipline.pipeline.sentinels import NotSpecified from zipline.pipeline.term import AssetExists, Slice from zipline.testing import parameter_space -from zipline.testing.predicates import assert_equal, assert_raises +from zipline.testing.predicates import ( + assert_equal, + assert_raises, + assert_raises_regex, + assert_regex, +) from zipline.utils.numpy_utils import ( bool_dtype, categorical_dtype, @@ -433,10 +440,50 @@ class ObjectIdentityTestCase(TestCase): with assert_raises(TypeError) as e: self.SomeFactorParameterized(a=[], b=[]) - assert_equal( + assert_regex( str(e.exception), - "SomeFactorParameterized expected a hashable value for parameter" - " 'a', but got [] instead.", + r"SomeFactorParameterized expected a hashable value for parameter" + r" '(a|b)', but got \[\] instead\.", + ) + + def test_parameterized_term_default_value(self): + defaults = {'a': 'default for a', 'b': 'default for b'} + + class F(Factor): + params = defaults + + inputs = (SomeDataSet.foo,) + dtype = 'f8' + window_length = 5 + + assert_equal(F().params, defaults) + assert_equal(F(a='new a').params, assoc(defaults, 'a', 'new a')) + assert_equal(F(b='new b').params, assoc(defaults, 'b', 'new b')) + assert_equal( + F(a='new a', b='new b').params, + {'a': 'new a', 'b': 'new b'}, + ) + + def test_parameterized_term_default_value_with_not_specified(self): + defaults = {'a': 'default for a', 'b': NotSpecified} + + class F(Factor): + params = defaults + + inputs = (SomeDataSet.foo,) + dtype = 'f8' + window_length = 5 + + pattern = r"F expected a keyword parameter 'b'\." + with assert_raises_regex(TypeError, pattern): + F() + with assert_raises_regex(TypeError, pattern): + F(a='new a') + + assert_equal(F(b='new b').params, assoc(defaults, 'b', 'new b')) + assert_equal( + F(a='new a', b='new b').params, + {'a': 'new a', 'b': 'new b'}, ) def test_bad_input(self): diff --git a/zipline/assets/symbol_rules.py b/zipline/assets/symbol_rules.py new file mode 100644 index 00000000..6a7fe0fc --- /dev/null +++ b/zipline/assets/symbol_rules.py @@ -0,0 +1,13 @@ +import re + + +_symbol_delimiter_regex = re.compile(r'[./\-_]') + + +def split_nasdaq(symbol): + sym = re.replace(_symbol_delimiter_regex, '', symbol) + return sym[:4], sym[4:] + + +def split_nyse(symbol): + return re.split(_symbol_delimiter_regex, symbol, maxsplit=1) diff --git a/zipline/pipeline/term.py b/zipline/pipeline/term.py index 563a0c81..4671daed 100644 --- a/zipline/pipeline/term.py +++ b/zipline/pipeline/term.py @@ -3,6 +3,7 @@ Base class for Filters, Factors and Classifiers """ from abc import ABCMeta, abstractproperty from bisect import insort +from collections import Mapping from weakref import WeakValueDictionary from numpy import ( @@ -146,10 +147,16 @@ class Term(with_metaclass(ABCMeta, object)): TypeError Raised if any parameter values are not passed or not hashable. """ + params = cls.params + if not isinstance(params, Mapping): + params = {k: NotSpecified for k in params} param_values = [] - for key in cls.params: + for key, default_value in params.items(): try: - value = kwargs.pop(key) + value = kwargs.pop(key, default_value) + if value is NotSpecified: + raise KeyError(key) + # Check here that the value is hashable so that we fail here # instead of trying to hash the param values tuple later. hash(value) @@ -171,8 +178,8 @@ class Term(with_metaclass(ABCMeta, object)): ) ) - param_values.append(value) - return tuple(zip(cls.params, param_values)) + param_values.append((key, value)) + return tuple(param_values) def __init__(self, *args, **kwargs): """ diff --git a/zipline/testing/predicates.py b/zipline/testing/predicates.py index b25167c3..e2f4078b 100644 --- a/zipline/testing/predicates.py +++ b/zipline/testing/predicates.py @@ -1,6 +1,8 @@ +from contextlib import contextmanager import datetime from functools import partial import inspect +import re from nose.tools import ( # noqa assert_almost_equal, @@ -205,6 +207,47 @@ def assert_is_subclass(subcls, cls, msg=''): ) +def assert_regex(result, expected, msg=''): + """Assert that ``expected`` matches the result. + + Parameters + ---------- + result : str + The string to search. + expected : str or compiled regex + The pattern to search for in ``result``. + msg : str, optional + An extra assertion message to print if this fails. + """ + assert re.search(expected, result), ( + '%s%r not found in %r' % (_fmt_msg(msg), expected, result) + ) + + +@contextmanager +def assert_raises_regex(exc, pattern, msg=''): + """Assert that some exception is raised in a context and that the message + matches some pattern. + + Parameters + ---------- + exc : type or tuple[type] + The exception type or types to expect. + pattern : str or compiled regex + The pattern to search for in the str of the raised exception. + msg : str, optional + An extra assertion message to print if this fails. + """ + try: + yield + except exc as e: + assert re.search(pattern, str(e)), ( + '%s%r not found in %r' % (_fmt_msg(msg), pattern, str(e)) + ) + else: + raise AssertionError('%s%s was not raised' % (_fmt_msg(msg), exc)) + + @dispatch(object, object) def assert_equal(result, expected, path=(), msg='', **kwargs): """Assert that two objects are equal using the ``==`` operator.