80 lines
2.4 KiB
Python
80 lines
2.4 KiB
Python
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
import pandas as pd
|
||
|
|
import pandas._testing as tm
|
||
|
|
|
||
|
|
|
||
|
|
def test_take(idx):
|
||
|
|
indexer = [4, 3, 0, 2]
|
||
|
|
result = idx.take(indexer)
|
||
|
|
expected = idx[indexer]
|
||
|
|
assert result.equals(expected)
|
||
|
|
|
||
|
|
# GH 10791
|
||
|
|
msg = "'MultiIndex' object has no attribute 'freq'"
|
||
|
|
with pytest.raises(AttributeError, match=msg):
|
||
|
|
idx.freq
|
||
|
|
|
||
|
|
|
||
|
|
def test_take_invalid_kwargs(idx):
|
||
|
|
idx = idx
|
||
|
|
indices = [1, 2]
|
||
|
|
|
||
|
|
msg = r"take\(\) got an unexpected keyword argument 'foo'"
|
||
|
|
with pytest.raises(TypeError, match=msg):
|
||
|
|
idx.take(indices, foo=2)
|
||
|
|
|
||
|
|
msg = "the 'out' parameter is not supported"
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
idx.take(indices, out=indices)
|
||
|
|
|
||
|
|
msg = "the 'mode' parameter is not supported"
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
idx.take(indices, mode="clip")
|
||
|
|
|
||
|
|
|
||
|
|
def test_take_fill_value():
|
||
|
|
# GH 12631
|
||
|
|
vals = [["A", "B"], [pd.Timestamp("2011-01-01"), pd.Timestamp("2011-01-02")]]
|
||
|
|
idx = pd.MultiIndex.from_product(vals, names=["str", "dt"])
|
||
|
|
|
||
|
|
result = idx.take(np.array([1, 0, -1]))
|
||
|
|
exp_vals = [
|
||
|
|
("A", pd.Timestamp("2011-01-02")),
|
||
|
|
("A", pd.Timestamp("2011-01-01")),
|
||
|
|
("B", pd.Timestamp("2011-01-02")),
|
||
|
|
]
|
||
|
|
expected = pd.MultiIndex.from_tuples(exp_vals, names=["str", "dt"])
|
||
|
|
tm.assert_index_equal(result, expected)
|
||
|
|
|
||
|
|
# fill_value
|
||
|
|
result = idx.take(np.array([1, 0, -1]), fill_value=True)
|
||
|
|
exp_vals = [
|
||
|
|
("A", pd.Timestamp("2011-01-02")),
|
||
|
|
("A", pd.Timestamp("2011-01-01")),
|
||
|
|
(np.nan, pd.NaT),
|
||
|
|
]
|
||
|
|
expected = pd.MultiIndex.from_tuples(exp_vals, names=["str", "dt"])
|
||
|
|
tm.assert_index_equal(result, expected)
|
||
|
|
|
||
|
|
# allow_fill=False
|
||
|
|
result = idx.take(np.array([1, 0, -1]), allow_fill=False, fill_value=True)
|
||
|
|
exp_vals = [
|
||
|
|
("A", pd.Timestamp("2011-01-02")),
|
||
|
|
("A", pd.Timestamp("2011-01-01")),
|
||
|
|
("B", pd.Timestamp("2011-01-02")),
|
||
|
|
]
|
||
|
|
expected = pd.MultiIndex.from_tuples(exp_vals, names=["str", "dt"])
|
||
|
|
tm.assert_index_equal(result, expected)
|
||
|
|
|
||
|
|
msg = "When allow_fill=True and fill_value is not None, all indices must be >= -1"
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
idx.take(np.array([1, 0, -2]), fill_value=True)
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
idx.take(np.array([1, 0, -5]), fill_value=True)
|
||
|
|
|
||
|
|
msg = "index -5 is out of bounds for( axis 0 with)? size 4"
|
||
|
|
with pytest.raises(IndexError, match=msg):
|
||
|
|
idx.take(np.array([1, -5]))
|