Merge pull request #287 from benjello/fix-gh-218

Fix #218
pydata · Sep 29, 2017 · 022d4a4 · 022d4a4
2 parents 6af5555 + 08cd904
commit 022d4a4
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/numexpr/expressions.py b/numexpr/expressions.py
@@ -216,8 +216,7 @@ def function(*args):
 @ophelper
 def where_func(a, b, c):
     if isinstance(a, ConstantNode):
-        #FIXME: This prevents where(True, a, b)
-        raise ValueError("too many dimensions")
+        return b if a.value else c
     if allConstantNodes([a, b, c]):
         return ConstantNode(numpy.where(a, b, c))
     return FuncNode('where', [a, b, c])

diff --git a/numexpr/tests/test_numexpr.py b/numexpr/tests/test_numexpr.py
@@ -295,6 +295,16 @@ def test_str_contains_long_needle(self):
         b = b'a' * 40
         res = evaluate('contains(a, b)')
         assert_equal(res, True)
+
+    def test_where_scalar_bool(self):
+        a = True
+        b = array([1, 2])
+        c = array([3, 4])
+        res = evaluate('where(a, b, c)')
+        assert_array_equal(res, b)
+        a = False
+        res = evaluate('where(a, b, c)')
+        assert_array_equal(res, c)
 
 
 class test_numexpr2(test_numexpr):