Fix deprecation warnings on numpy 1.10

njsmith · njsmith · commit d28e79bfabda · 2015-10-19T20:52:21.000-07:00
Apparently we were doing array == string and expecting a scalar to be
returned. That's not a great plan, so don't do that.
diff --git a/patsy/mgcv_cubic_splines.py b/patsy/mgcv_cubic_splines.py
@@ -10,7 +10,7 @@
 import numpy as np
 
 from patsy.util import (have_pandas, atleast_2d_column_default,
-                        no_pickling, assert_no_pickling)
+                        no_pickling, assert_no_pickling, safe_string_eq)
 from patsy.state import stateful_transform
 
 if have_pandas:
@@ -623,7 +623,7 @@ def memorize_finish(self):
         constraints = args["constraints"]
         n_constraints = 0
         if constraints is not None:
-            if constraints == "center":
+            if safe_string_eq(constraints, "center"):
                 # Here we collect only number of constraints,
                 # actual centering constraint will be computed after all_knots
                 n_constraints = 1
@@ -651,7 +651,7 @@ def memorize_finish(self):
                                                 lower_bound=args["lower_bound"],
                                                 upper_bound=args["upper_bound"])
         if constraints is not None:
-            if constraints == "center":
+            if safe_string_eq(constraints, "center"):
                 # Now we can compute centering constraints
                 constraints = _get_centering_constraint_from_dmatrix(
                     _get_free_crs_dmatrix(x, self._all_knots, cyclic=self._cyclic)
@@ -895,7 +895,7 @@ def __init__(self):
     def memorize_chunk(self, *args, **kwargs):
         constraints = self._tmp.setdefault("constraints",
                                            kwargs.get("constraints"))
-        if constraints == "center":
+        if safe_string_eq(constraints, "center"):
             args_2d = []
             for arg in args:
                 arg = atleast_2d_column_default(arg)
@@ -919,7 +919,7 @@ def memorize_finish(self):
         del self._tmp
 
         if constraints is not None:
-            if constraints == "center":
+            if safe_string_eq(constraints, "center"):
                 constraints = np.atleast_2d(tmp["sum"] / tmp["count"])
             else:
                 constraints = np.atleast_2d(constraints)
diff --git a/patsy/util.py b/patsy/util.py
@@ -20,6 +20,7 @@
            "safe_issubdtype",
            "no_pickling",
            "assert_no_pickling",
+           "safe_string_eq",
            ]
 
 import sys
@@ -699,3 +700,23 @@ def assert_no_pickling(obj):
     import pickle
     from nose.tools import assert_raises
     assert_raises(NotImplementedError, pickle.dumps, obj)
+
+# Use like:
+#   if safe_string_eq(constraints, "center"):
+#       ...
+# where 'constraints' might be a string or an array. (If it's an array, then
+# we can't use == becaues it might broadcast and ugh.)
+def safe_string_eq(obj, value):
+    if isinstance(obj, six.string_types):
+        return obj == value
+    else:
+        return False
+
+def test_safe_string_eq():
+    assert safe_string_eq("foo", "foo")
+    assert not safe_string_eq("foo", "bar")
+
+    if not six.PY3:
+        assert safe_string_eq(unicode("foo"), "foo")
+
+    assert not safe_string_eq(np.empty((2, 2)), "foo")
