From 5351b60a4c37aaf1ab19c7049298dc6f8d71d424 Mon Sep 17 00:00:00 2001
From: Joe Jevnik <joe@quantopian.com>
Date: Tue, 12 Jan 2016 17:30:56 -0500
Subject: [PATCH] ENH: adds optionally for preprocessors

---
 tests/utils/test_preprocess.py    | 20 +++++++++++++++
 zipline/utils/input_validation.py | 42 ++++++++++++++++++++++++++++++-
 2 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/tests/utils/test_preprocess.py b/tests/utils/test_preprocess.py
index 5270e605..4f425b16 100644
--- a/tests/utils/test_preprocess.py
+++ b/tests/utils/test_preprocess.py
@@ -17,6 +17,7 @@ from zipline.utils.input_validation import (
     expect_dtypes,
     expect_types,
     optional,
+    optionally,
 )
 
 
@@ -347,3 +348,22 @@ class PreprocessTestCase(TestCase):
         # test invalid timezone strings
         for tz in invalid:
             self.assertRaises(pytz.UnknownTimeZoneError, f, tz)
+
+    def test_optionally(self):
+        error = TypeError('arg must be int')
+
+        def preprocessor(func, argname, arg):
+            if not isinstance(arg, int):
+                raise error
+            return arg
+
+        @preprocess(a=optionally(preprocessor))
+        def f(a):
+            return a
+
+        self.assertIs(f(1), 1)
+        self.assertIsNone(f(None))
+
+        with self.assertRaises(TypeError) as e:
+            f('a')
+        self.assertIs(e.exception, error)
diff --git a/zipline/utils/input_validation.py b/zipline/utils/input_validation.py
index fc8f7c24..01544f59 100644
--- a/zipline/utils/input_validation.py
+++ b/zipline/utils/input_validation.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from datetime import tzinfo
-from functools import partial
+from functools import partial, wraps
 from operator import attrgetter
 
 from numpy import dtype
@@ -24,6 +24,46 @@ import toolz.curried.operator as op
 from zipline.utils.preprocess import preprocess
 
 
+def optionally(preprocessor):
+    """Modify a preprocessor to explicitly allow `None`.
+
+    Parameters
+    ----------
+    preprocessor : callable[callable, str, any -> any]
+        A preprocessor to delegate to when `arg is not None`.
+
+    Returns
+    -------
+    optional_preprocessor : callable[callable, str, any -> any]
+        A preprocessor that delegates to `preprocessor` when `arg is not None`.
+
+    Usage
+    -----
+    >>> def preprocessor(func, argname, arg):
+    ...     if not isinstance(arg, int):
+    ...         raise TypeError('arg must be int')
+    ...     return arg
+    ...
+    >>> @preprocess(a=optionally(preprocessor))
+    ... def f(a):
+    ...     return a
+    ...
+    >>> f(1)  # call with int
+    1
+    >> f('a')  # call with not int
+    Traceback (most recent call last):
+       ...
+    TypeError: arg must be int
+    >>> f(None)  # call with explicit None
+    None
+    """
+    @wraps(preprocessor)
+    def wrapper(func, argname, arg):
+        return arg if arg is None else preprocessor(func, argname, arg)
+
+    return wrapper
+
+
 def ensure_upper_case(func, argname, arg):
     if isinstance(arg, string_types):
         return arg.upper()