From 124555234098763427f7fe12de0fd87d8f8144ed Mon Sep 17 00:00:00 2001
From: Scott Sanderson <ssanderson@quantopian.com>
Date: Tue, 22 Mar 2016 13:30:43 -0400
Subject: [PATCH] DEV: Add expect_dimensions preprocessor.

---
 tests/utils/test_preprocess.py    | 38 ++++++++++++++++++++++++++-
 zipline/utils/input_validation.py | 43 +++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 1 deletion(-)

diff --git a/tests/utils/test_preprocess.py b/tests/utils/test_preprocess.py
index 2626cedb..343966ab 100644
--- a/tests/utils/test_preprocess.py
+++ b/tests/utils/test_preprocess.py
@@ -6,12 +6,13 @@ from types import FunctionType
 from unittest import TestCase
 
 from nose_parameterized import parameterized
-from numpy import arange, dtype
+from numpy import arange, array, dtype
 import pytz
 from six import PY3
 
 from zipline.utils.preprocess import call, preprocess
 from zipline.utils.input_validation import (
+    expect_dimensions,
     ensure_timezone,
     expect_element,
     expect_dtypes,
@@ -367,3 +368,38 @@ class PreprocessTestCase(TestCase):
         with self.assertRaises(TypeError) as e:
             f('a')
         self.assertIs(e.exception, error)
+
+    def test_expect_dimensions(self):
+
+        @expect_dimensions(x=2)
+        def foo(x, y):
+            return x[0, 0]
+
+        self.assertEqual(foo(arange(1).reshape(1, 1), 10), 0)
+
+        with self.assertRaises(ValueError) as e:
+            foo(arange(1), 1)
+        errmsg = str(e.exception)
+        expected = (
+            "{qualname}() expected a 2-D array for argument 'x', but got"
+            " a 1-D array instead.".format(qualname=qualname(foo))
+        )
+        self.assertEqual(errmsg, expected)
+
+        with self.assertRaises(ValueError) as e:
+            foo(arange(1).reshape(1, 1, 1), 1)
+        errmsg = str(e.exception)
+        expected = (
+            "{qualname}() expected a 2-D array for argument 'x', but got"
+            " a 3-D array instead.".format(qualname=qualname(foo))
+        )
+        self.assertEqual(errmsg, expected)
+
+        with self.assertRaises(ValueError) as e:
+            foo(array(0), 1)
+        errmsg = str(e.exception)
+        expected = (
+            "{qualname}() expected a 2-D array for argument 'x', but got"
+            " a scalar instead.".format(qualname=qualname(foo))
+        )
+        self.assertEqual(errmsg, expected)
diff --git a/zipline/utils/input_validation.py b/zipline/utils/input_validation.py
index 692cfaf7..5588518a 100644
--- a/zipline/utils/input_validation.py
+++ b/zipline/utils/input_validation.py
@@ -366,6 +366,49 @@ def expect_element(*_pos, **named):
     return preprocess(**valmap(_expect_element, named))
 
 
+def expect_dimensions(**dimensions):
+    """
+    Preprocessing decorator that verifies inputs are numpy arrays with a
+    specific dimensionality.
+
+    Usage
+    -----
+    >>> from numpy import array
+    >>> @expect_dimensions(x=1, y=2)
+    ... def foo(x, y):
+    ...    return x[0] + y[0, 0]
+    ...
+    >>> foo(array([1, 1]), array([[1, 1], [2, 2]]))
+    2
+    >>> foo(array([1, 1], array([1, 1])))
+    Traceback (most recent call last):
+       ...
+    TypeError: foo() expected a 2-D array for argument 'y', but got a 1-D array instead.  # noqa
+    """
+    def _expect_dimension(expected_ndim):
+        def _check(func, argname, argvalue):
+            funcname = _qualified_name(func)
+            actual_ndim = argvalue.ndim
+            if actual_ndim != expected_ndim:
+                if actual_ndim == 0:
+                    actual_repr = 'scalar'
+                else:
+                    actual_repr = "%d-D array" % actual_ndim
+                raise ValueError(
+                    "{func}() expected a {expected:d}-D array"
+                    " for argument {argname!r}, but got a {actual}"
+                    " instead.".format(
+                        func=funcname,
+                        expected=expected_ndim,
+                        argname=argname,
+                        actual=actual_repr,
+                    )
+                )
+            return argvalue
+        return _check
+    return preprocess(**valmap(_expect_dimension, dimensions))
+
+
 def coerce(from_, to, **to_kwargs):
     """
     A preprocessing decorator that coerces inputs of a given type by passing