test_xs.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  1. import numpy as np
  2. import pytest
  3. from pandas import (
  4. MultiIndex,
  5. Series,
  6. date_range,
  7. )
  8. import pandas._testing as tm
  9. def test_xs_datetimelike_wrapping():
  10. # GH#31630 a case where we shouldn't wrap datetime64 in Timestamp
  11. arr = date_range("2016-01-01", periods=3)._data._ndarray
  12. ser = Series(arr, dtype=object)
  13. for i in range(len(ser)):
  14. ser.iloc[i] = arr[i]
  15. assert ser.dtype == object
  16. assert isinstance(ser[0], np.datetime64)
  17. result = ser.xs(0)
  18. assert isinstance(result, np.datetime64)
  19. class TestXSWithMultiIndex:
  20. def test_xs_level_series(self, multiindex_dataframe_random_data):
  21. df = multiindex_dataframe_random_data
  22. ser = df["A"]
  23. expected = ser[:, "two"]
  24. result = df.xs("two", level=1)["A"]
  25. tm.assert_series_equal(result, expected)
  26. def test_series_getitem_multiindex_xs_by_label(self):
  27. # GH#5684
  28. idx = MultiIndex.from_tuples(
  29. [("a", "one"), ("a", "two"), ("b", "one"), ("b", "two")]
  30. )
  31. ser = Series([1, 2, 3, 4], index=idx)
  32. return_value = ser.index.set_names(["L1", "L2"], inplace=True)
  33. assert return_value is None
  34. expected = Series([1, 3], index=["a", "b"])
  35. return_value = expected.index.set_names(["L1"], inplace=True)
  36. assert return_value is None
  37. result = ser.xs("one", level="L2")
  38. tm.assert_series_equal(result, expected)
  39. def test_series_getitem_multiindex_xs(self):
  40. # GH#6258
  41. dt = list(date_range("20130903", periods=3))
  42. idx = MultiIndex.from_product([list("AB"), dt])
  43. ser = Series([1, 3, 4, 1, 3, 4], index=idx)
  44. expected = Series([1, 1], index=list("AB"))
  45. result = ser.xs("20130903", level=1)
  46. tm.assert_series_equal(result, expected)
  47. def test_series_xs_droplevel_false(self):
  48. # GH: 19056
  49. mi = MultiIndex.from_tuples(
  50. [("a", "x"), ("a", "y"), ("b", "x")], names=["level1", "level2"]
  51. )
  52. ser = Series([1, 1, 1], index=mi)
  53. result = ser.xs("a", axis=0, drop_level=False)
  54. expected = Series(
  55. [1, 1],
  56. index=MultiIndex.from_tuples(
  57. [("a", "x"), ("a", "y")], names=["level1", "level2"]
  58. ),
  59. )
  60. tm.assert_series_equal(result, expected)
  61. def test_xs_key_as_list(self):
  62. # GH#41760
  63. mi = MultiIndex.from_tuples([("a", "x")], names=["level1", "level2"])
  64. ser = Series([1], index=mi)
  65. with pytest.raises(TypeError, match="list keys are not supported"):
  66. ser.xs(["a", "x"], axis=0, drop_level=False)
  67. with pytest.raises(TypeError, match="list keys are not supported"):
  68. ser.xs(["a"], axis=0, drop_level=False)