test_truncate.py 2.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667
  1. from datetime import datetime
  2. import pytest
  3. import pandas as pd
  4. from pandas import (
  5. Series,
  6. date_range,
  7. )
  8. import pandas._testing as tm
  9. class TestTruncate:
  10. def test_truncate_datetimeindex_tz(self):
  11. # GH 9243
  12. idx = date_range("4/1/2005", "4/30/2005", freq="D", tz="US/Pacific")
  13. s = Series(range(len(idx)), index=idx)
  14. with pytest.raises(TypeError, match="Cannot compare tz-naive"):
  15. # GH#36148 as of 2.0 we require tzawareness compat
  16. s.truncate(datetime(2005, 4, 2), datetime(2005, 4, 4))
  17. lb = idx[1]
  18. ub = idx[3]
  19. result = s.truncate(lb.to_pydatetime(), ub.to_pydatetime())
  20. expected = Series([1, 2, 3], index=idx[1:4])
  21. tm.assert_series_equal(result, expected)
  22. def test_truncate_periodindex(self):
  23. # GH 17717
  24. idx1 = pd.PeriodIndex(
  25. [pd.Period("2017-09-02"), pd.Period("2017-09-02"), pd.Period("2017-09-03")]
  26. )
  27. series1 = Series([1, 2, 3], index=idx1)
  28. result1 = series1.truncate(after="2017-09-02")
  29. expected_idx1 = pd.PeriodIndex(
  30. [pd.Period("2017-09-02"), pd.Period("2017-09-02")]
  31. )
  32. tm.assert_series_equal(result1, Series([1, 2], index=expected_idx1))
  33. idx2 = pd.PeriodIndex(
  34. [pd.Period("2017-09-03"), pd.Period("2017-09-02"), pd.Period("2017-09-03")]
  35. )
  36. series2 = Series([1, 2, 3], index=idx2)
  37. result2 = series2.sort_index().truncate(after="2017-09-02")
  38. expected_idx2 = pd.PeriodIndex([pd.Period("2017-09-02")])
  39. tm.assert_series_equal(result2, Series([2], index=expected_idx2))
  40. def test_truncate_one_element_series(self):
  41. # GH 35544
  42. series = Series([0.1], index=pd.DatetimeIndex(["2020-08-04"]))
  43. before = pd.Timestamp("2020-08-02")
  44. after = pd.Timestamp("2020-08-04")
  45. result = series.truncate(before=before, after=after)
  46. # the input Series and the expected Series are the same
  47. tm.assert_series_equal(result, series)
  48. def test_truncate_index_only_one_unique_value(self):
  49. # GH 42365
  50. obj = Series(0, index=date_range("2021-06-30", "2021-06-30")).repeat(5)
  51. truncated = obj.truncate("2021-06-28", "2021-07-01")
  52. tm.assert_series_equal(truncated, obj)