Fixed bug in numpy.piecewise()

ericsuh · ericsuh · commit e553e1b4595a · 2012-07-05T16:57:06.000-03:00
Updated numpy/lib/function_base.py to fix bug in numpy.piecewise() for 0-d
array handling and boolean indexing with scalars.
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
@@ -673,21 +673,11 @@ def piecewise(x, condlist, funclist, *args, **kw):
     """
     x = asanyarray(x)
     n2 = len(funclist)
-    if isscalar(condlist) or \
-           not (isinstance(condlist[0], list) or
-                isinstance(condlist[0], ndarray)):
+    if isscalar(condlist):
         condlist = [condlist]
     condlist = [asarray(c, dtype=bool) for c in condlist]
     n = len(condlist)
-    if n == n2-1:  # compute the "otherwise" condition.
-        totlist = condlist[0]
-        for k in range(1, n):
-            totlist |= condlist[k]
-        condlist.append(~totlist)
-        n += 1
-    if (n != n2):
-        raise ValueError(
-                "function list and condition list must be the same")
+
     zerod = False
     # This is a hack to work around problems with NumPy's
     #  handling of 0-d arrays and boolean indexing with
@@ -704,6 +694,16 @@ def piecewise(x, condlist, funclist, *args, **kw):
             newcondlist.append(condition)
         condlist = newcondlist
 
+    if n == n2-1:  # compute the "otherwise" condition.
+        totlist = condlist[0]
+        for k in range(1, n):
+            totlist |= condlist[k]
+        condlist.append(~totlist)
+        n += 1
+    if (n != n2):
+        raise ValueError(
+                "function list and condition list must be the same")
+
     y = zeros(x.shape, x.dtype)
     for k in range(n):
         item = funclist[k]
