Fix piecewise to handle 0-d inputs.

2 files changed, 27 insertions, 1 deletions

diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index d2bb47552..cb2cbf29b 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -574,13 +574,32 @@ def piecewise(x, condlist, funclist, *args, **kw):
         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 
+    #  numpy.bool_ scalars
+    if x.ndim == 0:
+        x = x[None]
+        zerod = True
+        newcondlist = []
+        for k in range(n):
+            if condlist[k].ndim == 0:
+                condition = condlist[k][None]
+            else:
+                condition = condlist[k]
+            newcondlist.append(condition)
+        condlist = newcondlist
     y = empty(x.shape, x.dtype)
     for k in range(n):
         item = funclist[k]
         if not callable(item):
             y[condlist[k]] = item
         else:
-            y[condlist[k]] = item(x[condlist[k]], *args, **kw)
+            vals = x[condlist[k]]
+            if vals.size > 0:
+                y[condlist[k]] = item(vals, *args, **kw)
+    if zerod:
+        y = y.squeeze()
     return y
 
 def select(condlist, choicelist, default=0):
diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index 3cd7cd6bb..c885f49af 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -616,5 +616,12 @@ def compare_results(res,desired):
     for i in range(len(desired)):
         assert_array_equal(res[i],desired[i])
 
+class TestPiecewise(TestCase):
+    def test_0d(self):
+        x = array(3)
+        y = piecewise(x, x>3, [4, 0])
+        assert y.ndim == 0
+        assert y == 0
+
 if __name__ == "__main__":
     nose.run(argv=['', __file__])