Doc and move tests

TomAugspurger · TomAugspurger · commit 5db66248f2c4 · 2018-04-26T07:17:30.000-05:00
diff --git a/pandas/core/algorithms.py b/pandas/core/algorithms.py
@@ -1455,19 +1455,21 @@ def take(arr, indexer, allow_fill=False, fill_value=None):
 
     Parameters
     ----------
-    arr : ndarray or ExtensionArray
+    arr : sequence
+        Non array-likes (sequences without a dtype) are coereced
+        to an ndarray.
     indexer : sequence of integers
-        Indices to be taken. See Notes for how negative indicies
-        are handled.
+        Indices to be taken.
     allow_fill : bool, default False
         How to handle negative values in `indexer`.
 
-        For False values (the default), negative values in `indexer`
-        indiciate slices from the right.
+        * False: negative values in `indexer` indicate
+          slices from the right (the default)
+
+        * True: negative values in `indexer` indicate
+          missing values. These values are set to `fill_value`. Any other
+          other negative values raise a ``ValueError``.
 
-        For True values, indicies where `indexer` is ``-1`` indicate
-        missing values. These values are set to `fill_value`. Any other
-        other negative value should raise a ``ValueError``.
     fill_value : any, optional
         Fill value to use for NA-indicies when `allow_fill` is True.
         This may be ``None``, in which case the default NA value for
@@ -1486,13 +1488,42 @@ def take(arr, indexer, allow_fill=False, fill_value=None):
         When the indexer contains negative values other than ``-1``
         and `allow_fill` is True.
 
+    Notes
+    -----
+    When `allow_fill` is False, `indexer` may be whatever dimensionality
+    is accepted by NumPy for `arr`.
+
+    When `allow_fill` is True, `indexer` should be 1-D.
+
     See Also
     --------
     numpy.take
+
+    Examples
+    --------
+    >>> from pandas.api.extensions import take
+
+    With the default ``allow_fill=False``, negative numbers indicate
+    slices from the right.
+
+    >>> take(np.array([10, 20, 30]), [0, 0, -1])
+    array([10, 10, 30])
+
+    Setting ``allow_fill=True`` will place `fill_value` in those positions.
+
+    >>> take(np.array([10, 20, 30]), [0, 0, -1], allow_fill=True)
+    array([10., 10., nan])
+
+    >>> take(np.array([10, 20, 30]), [0, 0, -1], allow_fill=True,
+    ...      fill_value=-10)
+    array([ 10,  10, -10])
     """
     from pandas.core.indexing import validate_indices
 
-    # Do we require int64 here?
+    if not is_array_like(arr):
+        arr = np.asarray(arr)
+
+    # Do we require int64 or intp here?
     indexer = np.asarray(indexer, dtype='int')
 
     if allow_fill:
diff --git a/pandas/tests/test_algos.py b/pandas/tests/test_algos.py
@@ -1564,40 +1564,3 @@ def test_index(self):
         idx = Index(['1 day', '1 day', '-1 day', '-1 day 2 min',
                      '2 min', '2 min'], dtype='timedelta64[ns]')
         tm.assert_series_equal(algos.mode(idx), exp)
-
-
-class TestTake(object):
-
-    def test_bounds_check_large(self):
-        arr = np.array([1, 2])
-        with pytest.raises(IndexError):
-            algos.take(arr, [2, 3], allow_fill=True)
-
-        with pytest.raises(IndexError):
-            algos.take(arr, [2, 3], allow_fill=False)
-
-    def test_bounds_check_small(self):
-        arr = np.array([1, 2, 3], dtype=np.int64)
-        indexer = [0, -1, -2]
-        with pytest.raises(ValueError):
-            algos.take(arr, indexer, allow_fill=True)
-
-        result = algos.take(arr, indexer)
-        expected = np.array([1, 3, 2], dtype=np.int64)
-        tm.assert_numpy_array_equal(result, expected)
-
-    @pytest.mark.parametrize('allow_fill', [True, False])
-    def test_take_empty(self, allow_fill):
-        arr = np.array([], dtype=np.int64)
-        # empty take is ok
-        result = algos.take(arr, [], allow_fill=allow_fill)
-        tm.assert_numpy_array_equal(arr, result)
-
-        with pytest.raises(IndexError):
-            algos.take(arr, [0], allow_fill=allow_fill)
-
-    def test_take_na_empty(self):
-        result = algos.take(np.array([]), [-1, -1], allow_fill=True,
-                            fill_value=0.0)
-        expected = np.array([0., 0.])
-        tm.assert_numpy_array_equal(result, expected)
diff --git a/pandas/tests/test_take.py b/pandas/tests/test_take.py
@@ -3,6 +3,7 @@
 from datetime import datetime
 
 import numpy as np
+import pytest
 from pandas.compat import long
 import pandas.core.algorithms as algos
 import pandas.util.testing as tm
@@ -445,3 +446,47 @@ def test_2d_datetime64(self):
         expected = arr.take(indexer, axis=1)
         expected[:, [2, 4]] = datetime(2007, 1, 1)
         tm.assert_almost_equal(result, expected)
+
+
+class TestExtensionTake(object):
+    # The take method found in pd.api.extensions
+
+    def test_bounds_check_large(self):
+        arr = np.array([1, 2])
+        with pytest.raises(IndexError):
+            algos.take(arr, [2, 3], allow_fill=True)
+
+        with pytest.raises(IndexError):
+            algos.take(arr, [2, 3], allow_fill=False)
+
+    def test_bounds_check_small(self):
+        arr = np.array([1, 2, 3], dtype=np.int64)
+        indexer = [0, -1, -2]
+        with pytest.raises(ValueError):
+            algos.take(arr, indexer, allow_fill=True)
+
+        result = algos.take(arr, indexer)
+        expected = np.array([1, 3, 2], dtype=np.int64)
+        tm.assert_numpy_array_equal(result, expected)
+
+    @pytest.mark.parametrize('allow_fill', [True, False])
+    def test_take_empty(self, allow_fill):
+        arr = np.array([], dtype=np.int64)
+        # empty take is ok
+        result = algos.take(arr, [], allow_fill=allow_fill)
+        tm.assert_numpy_array_equal(arr, result)
+
+        with pytest.raises(IndexError):
+            algos.take(arr, [0], allow_fill=allow_fill)
+
+    def test_take_na_empty(self):
+        result = algos.take(np.array([]), [-1, -1], allow_fill=True,
+                            fill_value=0.0)
+        expected = np.array([0., 0.])
+        tm.assert_numpy_array_equal(result, expected)
+
+    def test_take_coerces_list(self):
+        arr = [1, 2, 3]
+        result = algos.take(arr, [0, 0])
+        expected = np.array([1, 1])
+        tm.assert_numpy_array_equal(result, expected)
