Skip to content

Commit 43c8313

Browse files
hannarojuliantaylor
authored andcommitted
BUG: Fixes numpy#5524 and adds test
argpartition does not fail anymore on non-ndarray array-likes. Fix as implemented by @maniteja123.
1 parent 1f72899 commit 43c8313

File tree

2 files changed

+15
-1
lines changed

2 files changed

+15
-1
lines changed

numpy/core/fromnumeric.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -679,8 +679,16 @@ def argpartition(a, kth, axis=-1, kind='introselect', order=None):
679679
>>> x[np.argpartition(x, (1, 3))]
680680
array([1, 2, 3, 4])
681681
682+
>>> x = [3, 4, 2, 1]
683+
>>> np.array(x)[np.argpartition(x, 3)]
684+
array([2, 1, 3, 4])
685+
682686
"""
683-
return a.argpartition(kth, axis, kind=kind, order=order)
687+
try:
688+
argpartition = a.argpartition
689+
except AttributeError:
690+
return _wrapit(a, 'argpartition',kth, axis, kind, order)
691+
return argpartition(kth, axis, kind=kind, order=order)
684692

685693

686694
def sort(a, axis=-1, kind='quicksort', order=None):

numpy/core/tests/test_multiarray.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1702,6 +1702,12 @@ def test_partition_fuzz(self):
17021702
assert_array_equal(np.partition(d, kth)[kth], tgt,
17031703
err_msg="data: %r\n kth: %r" % (d, kth))
17041704

1705+
def test_argpartition_gh5524(self):
1706+
# A test for functionality of argpartition on lists.
1707+
d = [6,7,3,2,9,0]
1708+
p = np.argpartition(d,1)
1709+
self.assert_partitioned(np.array(d)[p],[1])
1710+
17051711
def test_flatten(self):
17061712
x0 = np.array([[1, 2, 3], [4, 5, 6]], np.int32)
17071713
x1 = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], np.int32)

0 commit comments

Comments
 (0)