test_indexing.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  1. import numpy as np
  2. import pytest
  3. import pandas as pd
  4. from pandas import (
  5. Index,
  6. NaT,
  7. )
  8. import pandas._testing as tm
  9. class TestGetSliceBounds:
  10. @pytest.mark.parametrize("side, expected", [("left", 4), ("right", 5)])
  11. def test_get_slice_bounds_within(self, side, expected):
  12. index = Index(list("abcdef"))
  13. result = index.get_slice_bound("e", side=side)
  14. assert result == expected
  15. @pytest.mark.parametrize("side", ["left", "right"])
  16. @pytest.mark.parametrize(
  17. "data, bound, expected", [(list("abcdef"), "x", 6), (list("bcdefg"), "a", 0)]
  18. )
  19. def test_get_slice_bounds_outside(self, side, expected, data, bound):
  20. index = Index(data)
  21. result = index.get_slice_bound(bound, side=side)
  22. assert result == expected
  23. def test_get_slice_bounds_invalid_side(self):
  24. with pytest.raises(ValueError, match="Invalid value for side kwarg"):
  25. Index([]).get_slice_bound("a", side="middle")
  26. class TestGetIndexerNonUnique:
  27. def test_get_indexer_non_unique_dtype_mismatch(self):
  28. # GH#25459
  29. indexes, missing = Index(["A", "B"]).get_indexer_non_unique(Index([0]))
  30. tm.assert_numpy_array_equal(np.array([-1], dtype=np.intp), indexes)
  31. tm.assert_numpy_array_equal(np.array([0], dtype=np.intp), missing)
  32. class TestGetLoc:
  33. @pytest.mark.slow # to_flat_index takes a while
  34. def test_get_loc_tuple_monotonic_above_size_cutoff(self):
  35. # Go through the libindex path for which using
  36. # _bin_search vs ndarray.searchsorted makes a difference
  37. lev = list("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
  38. dti = pd.date_range("2016-01-01", periods=100)
  39. mi = pd.MultiIndex.from_product([lev, range(10**3), dti])
  40. oidx = mi.to_flat_index()
  41. loc = len(oidx) // 2
  42. tup = oidx[loc]
  43. res = oidx.get_loc(tup)
  44. assert res == loc
  45. def test_get_loc_nan_object_dtype_nonmonotonic_nonunique(self):
  46. # case that goes through _maybe_get_bool_indexer
  47. idx = Index(["foo", np.nan, None, "foo", 1.0, None], dtype=object)
  48. # we dont raise KeyError on nan
  49. res = idx.get_loc(np.nan)
  50. assert res == 1
  51. # we only match on None, not on np.nan
  52. res = idx.get_loc(None)
  53. expected = np.array([False, False, True, False, False, True])
  54. tm.assert_numpy_array_equal(res, expected)
  55. # we don't match at all on mismatched NA
  56. with pytest.raises(KeyError, match="NaT"):
  57. idx.get_loc(NaT)
  58. def test_getitem_boolean_ea_indexer():
  59. # GH#45806
  60. ser = pd.Series([True, False, pd.NA], dtype="boolean")
  61. result = ser.index[ser]
  62. expected = Index([0])
  63. tm.assert_index_equal(result, expected)