From 22ffe3fe49159290c82ae7ca6e088545336ac9d4 Mon Sep 17 00:00:00 2001
From: llllllllll <joe@quantopian.com>
Date: Mon, 12 Oct 2015 19:15:27 -0400
Subject: [PATCH] ENH: adds expect_element preprocessor

---
 tests/utils/test_preprocess.py    | 23 +++++++++-
 zipline/utils/input_validation.py | 74 +++++++++++++++++++++++++------
 2 files changed, 82 insertions(+), 15 deletions(-)

diff --git a/tests/utils/test_preprocess.py b/tests/utils/test_preprocess.py
index 4ed686fb..a9728567 100644
--- a/tests/utils/test_preprocess.py
+++ b/tests/utils/test_preprocess.py
@@ -6,7 +6,11 @@ from unittest import TestCase
 from nose_parameterized import parameterized
 
 from zipline.utils.preprocess import call, preprocess
-from zipline.utils.input_validation import expect_types, optional
+from zipline.utils.input_validation import (
+    expect_element,
+    expect_types,
+    optional,
+)
 
 
 def noop(func, argname, argvalue):
@@ -225,3 +229,20 @@ class PreprocessTestCase(TestCase):
             "type int or NoneType for argument 'a', but got str instead."
         ).format(modname=foo.__module__)
         self.assertEqual(e.exception.args[0], expected_message)
+
+    def test_expect_element(self):
+        @expect_element(a={'a', 'b'})
+        def f(a):
+            return a
+
+        self.assertEqual(f('a'), 'a')
+        self.assertEqual(f('b'), 'b')
+
+        with self.assertRaises(ValueError) as e:
+            f('c')
+
+        expected_message = (
+            "{modname}.f() expected a value in {{'a', 'b'}}"
+            " for argument 'a', but got 'c' instead."
+        ).format(modname=f.__module__)
+        self.assertEqual(e.exception.args[0], expected_message)
diff --git a/zipline/utils/input_validation.py b/zipline/utils/input_validation.py
index 9ee0e109..e847b3d3 100644
--- a/zipline/utils/input_validation.py
+++ b/zipline/utils/input_validation.py
@@ -13,7 +13,8 @@
 # limitations under the License.
 
 from six import iteritems, string_types
-from toolz import valmap
+from toolz import valmap, complement, flip, compose
+import toolz.curried.operator as op
 
 from zipline.utils.preprocess import preprocess
 
@@ -71,14 +72,28 @@ def _qualified_name(obj):
     return '.'.join([module, obj.__name__])
 
 
+def _mk_check(exc, template, pred, actual):
+    def _check(func, argname, argvalue):
+        if pred(argvalue):
+            raise exc(
+                template % {
+                    'funcname': _qualified_name(func),
+                    'argname': argname,
+                    'actual': actual(argvalue),
+                },
+            )
+        return argvalue
+    return _check
+
+
 def _expect_type(type_):
     """
     Factory for type-checking functions that work the @preprocess decorator.
     """
     # Slightly different messages for type and tuple of types.
     _template = (
-        "{{funcname}}() expected a value of type {type_or_types} "
-        "for argument '{{argname}}', but got {{actual}} instead."
+        "%(funcname)s() expected a value of type {type_or_types} "
+        "for argument '%(argname)s', but got %(actual)s instead."
     )
     if isinstance(type_, tuple):
         template = _template.format(
@@ -87,17 +102,12 @@ def _expect_type(type_):
     else:
         template = _template.format(type_or_types=_qualified_name(type_))
 
-    def _check_type(func, argname, argvalue):
-        if not isinstance(argvalue, type_):
-            raise TypeError(
-                template.format(
-                    funcname=_qualified_name(func),
-                    argname=argname,
-                    actual=_qualified_name(type(argvalue)),
-                )
-            )
-        return argvalue
-    return _check_type
+    return _mk_check(
+        TypeError,
+        template,
+        complement(flip(isinstance, type_)),
+        compose(_qualified_name, type),
+    )
 
 
 def optional(type_):
@@ -122,3 +132,39 @@ def optional(type_):
     False
     """
     return (type_, type(None))
+
+
+def expect_element(*_pos, **named):
+    """
+    Preprocessing decorator that verifies inputs have expected types.
+
+    Usage
+    -----
+    >>> @expect_types(x=int, y=str)
+    ... def foo(x, y):
+    ...    return x, y
+    ...
+    >>> foo(2, '3')
+    (2, '3')
+    >>> foo(2.0, '3')
+    Traceback (most recent call last):
+       ...
+    TypeError: foo() expected an argument of type 'int' for argument 'x', but got float instead.  # noqa
+    """
+    if _pos:
+        raise TypeError("expect_element() only takes keyword arguments.")
+
+    return preprocess(**valmap(_expect_element, named))
+
+
+def _expect_element(collection):
+    template = (
+        "%(funcname)s() expected a value in {collection} "
+        "for argument '%(argname)s', but got %(actual)s instead."
+    ).format(collection=collection)
+    return _mk_check(
+        ValueError,
+        template,
+        complement(op.contains(collection)),
+        repr,
+    )