test_take.py 963 B

123456789101112131415161718192021222324252627282930313233
  1. import pytest
  2. import pandas as pd
  3. from pandas import Series
  4. import pandas._testing as tm
  5. def test_take():
  6. ser = Series([-1, 5, 6, 2, 4])
  7. actual = ser.take([1, 3, 4])
  8. expected = Series([5, 2, 4], index=[1, 3, 4])
  9. tm.assert_series_equal(actual, expected)
  10. actual = ser.take([-1, 3, 4])
  11. expected = Series([4, 2, 4], index=[4, 3, 4])
  12. tm.assert_series_equal(actual, expected)
  13. msg = lambda x: f"index {x} is out of bounds for( axis 0 with)? size 5"
  14. with pytest.raises(IndexError, match=msg(10)):
  15. ser.take([1, 10])
  16. with pytest.raises(IndexError, match=msg(5)):
  17. ser.take([2, 5])
  18. def test_take_categorical():
  19. # https://github.com/pandas-dev/pandas/issues/20664
  20. ser = Series(pd.Categorical(["a", "b", "c"]))
  21. result = ser.take([-2, -2, 0])
  22. expected = Series(
  23. pd.Categorical(["b", "b", "a"], categories=["a", "b", "c"]), index=[1, 1, 0]
  24. )
  25. tm.assert_series_equal(result, expected)