test_subclass.py 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. import numpy as np
  2. import pytest
  3. import pandas as pd
  4. import pandas._testing as tm
  5. class TestSeriesSubclassing:
  6. @pytest.mark.parametrize(
  7. "idx_method, indexer, exp_data, exp_idx",
  8. [
  9. ["loc", ["a", "b"], [1, 2], "ab"],
  10. ["iloc", [2, 3], [3, 4], "cd"],
  11. ],
  12. )
  13. def test_indexing_sliced(self, idx_method, indexer, exp_data, exp_idx):
  14. s = tm.SubclassedSeries([1, 2, 3, 4], index=list("abcd"))
  15. res = getattr(s, idx_method)[indexer]
  16. exp = tm.SubclassedSeries(exp_data, index=list(exp_idx))
  17. tm.assert_series_equal(res, exp)
  18. def test_to_frame(self):
  19. s = tm.SubclassedSeries([1, 2, 3, 4], index=list("abcd"), name="xxx")
  20. res = s.to_frame()
  21. exp = tm.SubclassedDataFrame({"xxx": [1, 2, 3, 4]}, index=list("abcd"))
  22. tm.assert_frame_equal(res, exp)
  23. def test_subclass_unstack(self):
  24. # GH 15564
  25. s = tm.SubclassedSeries([1, 2, 3, 4], index=[list("aabb"), list("xyxy")])
  26. res = s.unstack()
  27. exp = tm.SubclassedDataFrame({"x": [1, 3], "y": [2, 4]}, index=["a", "b"])
  28. tm.assert_frame_equal(res, exp)
  29. def test_subclass_empty_repr(self):
  30. sub_series = tm.SubclassedSeries()
  31. assert "SubclassedSeries" in repr(sub_series)
  32. def test_asof(self):
  33. N = 3
  34. rng = pd.date_range("1/1/1990", periods=N, freq="53s")
  35. s = tm.SubclassedSeries({"A": [np.nan, np.nan, np.nan]}, index=rng)
  36. result = s.asof(rng[-2:])
  37. assert isinstance(result, tm.SubclassedSeries)
  38. def test_explode(self):
  39. s = tm.SubclassedSeries([[1, 2, 3], "foo", [], [3, 4]])
  40. result = s.explode()
  41. assert isinstance(result, tm.SubclassedSeries)
  42. def test_equals(self):
  43. # https://github.com/pandas-dev/pandas/pull/34402
  44. # allow subclass in both directions
  45. s1 = pd.Series([1, 2, 3])
  46. s2 = tm.SubclassedSeries([1, 2, 3])
  47. assert s1.equals(s2)
  48. assert s2.equals(s1)
  49. class SubclassedSeries(pd.Series):
  50. @property
  51. def _constructor(self):
  52. def _new(*args, **kwargs):
  53. # some constructor logic that accesses the Series' name
  54. if self.name == "test":
  55. return pd.Series(*args, **kwargs)
  56. return SubclassedSeries(*args, **kwargs)
  57. return _new
  58. def test_constructor_from_dict():
  59. # https://github.com/pandas-dev/pandas/issues/52445
  60. result = SubclassedSeries({"a": 1, "b": 2, "c": 3})
  61. assert isinstance(result, SubclassedSeries)