Author: Sergey Matyunin <sbmatyu...@gmail.com>
Branch: where_1_arg
Changeset: r83020:35aba6438e37
Date: 2016-03-13 19:43 +0100
http://bitbucket.org/pypy/pypy/changeset/35aba6438e37/
Log: Implemented where for 1 argument. Added tests.
diff --git a/pypy/module/micronumpy/arrayops.py
b/pypy/module/micronumpy/arrayops.py
--- a/pypy/module/micronumpy/arrayops.py
+++ b/pypy/module/micronumpy/arrayops.py
@@ -71,8 +71,8 @@
"""
if space.is_none(w_y):
if space.is_none(w_x):
- raise OperationError(space.w_NotImplementedError, space.wrap(
- "1-arg where unsupported right now"))
+ arr = convert_to_array(space, w_arr)
+ return arr.descr_nonzero(space)
raise OperationError(space.w_ValueError, space.wrap(
"Where should be called with either 1 or 3 arguments"))
if space.is_none(w_x):
diff --git a/pypy/module/micronumpy/test/test_arrayops.py
b/pypy/module/micronumpy/test/test_arrayops.py
--- a/pypy/module/micronumpy/test/test_arrayops.py
+++ b/pypy/module/micronumpy/test/test_arrayops.py
@@ -54,8 +54,24 @@
assert (where(False, 1, [1, 2, 3]) == [1, 2, 3]).all()
assert (where([1, 2, 3], True, False) == [True, True, True]).all()
- #def test_where_1_arg(self):
- # xxx
+ def test_where_1_arg(self):
+ from numpy import where, array
+
+ result = where([1,0,1])
+
+ assert isinstance(result, tuple)
+ assert len(result) == 1
+ assert (result[0] == array([0, 2])).all()
+
+ def test_where_1_arg_2d(self):
+ from numpy import where, array
+
+ result = where([[1,0,1],[2,-1,-1]])
+
+ assert isinstance(result, tuple)
+ assert len(result) == 2
+ assert (result[0] == array([0, 0, 1, 1, 1])).all()
+ assert (result[1] == array([0, 2, 0, 1, 2])).all()
def test_where_invalidates(self):
from numpy import where, ones, zeros, array
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
https://mail.python.org/mailman/listinfo/pypy-commit
