Refactor: push eval_env error checking into EvalEnvironment.capture where it belongs

Author: njsmith

patsy/build.py

@@ -615,7 +615,8 @@ def _make_term_column_builders(terms,
             term_to_column_builders[term] = column_builders
     return new_term_order, term_to_column_builders
 
-def design_matrix_builders(termlists, data_iter_maker, eval_env, NA_action="drop"):
+def design_matrix_builders(termlists, data_iter_maker, eval_env,
+                           NA_action="drop"):
     """Construct several :class:`DesignMatrixBuilders` from termlists.
 
     This is one of Patsy's fundamental functions. This function and
@@ -654,14 +655,10 @@ def design_matrix_builders(termlists, data_iter_maker, eval_env, NA_action="drop
     .. versionadded:: 0.4.0
        The ``eval_env`` argument.
     """
-    # Check type of eval_env to help people migrating to 0.4.0. Third
-    # argument used to be NA_action (a string). Having the check for
-    # eval_env's type gives people migrating to 0.4.0 who used NA_action
-    # not as a keyword argument a nice error message here, instead of a
-    # more obscure backtrace later on.
-    if not isinstance(eval_env, six.integer_types + (EvalEnvironment,)):
-        raise TypeError("Parameter 'eval_env' must be either an integer or an instance "
-                        "of patsy.EvalEnvironment.")
+    # People upgrading from versions prior to 0.4.0 could potentially have
+    # passed NA_action as the 3rd positional argument. Fortunately
+    # EvalEnvironment.capture only accepts int and EvalEnvironment objects,
+    # and we improved its error messages to make this clear.
     eval_env = EvalEnvironment.capture(eval_env, reference=1)
     if isinstance(NA_action, str):
         NA_action = NAAction(NA_action)

patsy/eval.py

@@ -16,8 +16,9 @@
 import __future__
 import inspect
 import tokenize
-import six
 import ast
+import numbers
+import six
 from patsy import PatsyError
 from patsy.util import PushbackAdapter
 from patsy.tokens import (pretty_untokenize, normalize_token_spacing,
@@ -213,8 +214,11 @@ def my_model(formula_like, data, eval_env=0):
         """
         if isinstance(eval_env, cls):
             return eval_env
-        else:
+        elif isinstance(eval_env, numbers.Integral):
             depth = eval_env + reference
+        else:
+            raise TypeError("Parameter 'eval_env' must be either an integer "
+                            "or an instance of patsy.EvalEnvironment.")
         frame = inspect.currentframe()
         try:
             for i in range(depth + 1):
@@ -293,6 +297,8 @@ def test_EvalEnvironment_capture_namespace():
 
     assert EvalEnvironment.capture(b1) is b1
 
+    assert_raises(TypeError, EvalEnvironment.capture, 1.2)
+
 def test_EvalEnvironment_capture_flags():
     if sys.version_info >= (3,):
         # This is the only __future__ feature currently usable in Python