test_clip.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138
  1. from datetime import datetime
  2. import numpy as np
  3. import pytest
  4. import pandas as pd
  5. from pandas import (
  6. Series,
  7. Timestamp,
  8. isna,
  9. notna,
  10. )
  11. import pandas._testing as tm
  12. class TestSeriesClip:
  13. def test_clip(self, datetime_series):
  14. val = datetime_series.median()
  15. assert datetime_series.clip(lower=val).min() == val
  16. assert datetime_series.clip(upper=val).max() == val
  17. result = datetime_series.clip(-0.5, 0.5)
  18. expected = np.clip(datetime_series, -0.5, 0.5)
  19. tm.assert_series_equal(result, expected)
  20. assert isinstance(expected, Series)
  21. def test_clip_types_and_nulls(self):
  22. sers = [
  23. Series([np.nan, 1.0, 2.0, 3.0]),
  24. Series([None, "a", "b", "c"]),
  25. Series(pd.to_datetime([np.nan, 1, 2, 3], unit="D")),
  26. ]
  27. for s in sers:
  28. thresh = s[2]
  29. lower = s.clip(lower=thresh)
  30. upper = s.clip(upper=thresh)
  31. assert lower[notna(lower)].min() == thresh
  32. assert upper[notna(upper)].max() == thresh
  33. assert list(isna(s)) == list(isna(lower))
  34. assert list(isna(s)) == list(isna(upper))
  35. def test_series_clipping_with_na_values(self, any_numeric_ea_dtype, nulls_fixture):
  36. # Ensure that clipping method can handle NA values with out failing
  37. # GH#40581
  38. if nulls_fixture is pd.NaT:
  39. # constructor will raise, see
  40. # test_constructor_mismatched_null_nullable_dtype
  41. return
  42. ser = Series([nulls_fixture, 1.0, 3.0], dtype=any_numeric_ea_dtype)
  43. s_clipped_upper = ser.clip(upper=2.0)
  44. s_clipped_lower = ser.clip(lower=2.0)
  45. expected_upper = Series([nulls_fixture, 1.0, 2.0], dtype=any_numeric_ea_dtype)
  46. expected_lower = Series([nulls_fixture, 2.0, 3.0], dtype=any_numeric_ea_dtype)
  47. tm.assert_series_equal(s_clipped_upper, expected_upper)
  48. tm.assert_series_equal(s_clipped_lower, expected_lower)
  49. def test_clip_with_na_args(self):
  50. """Should process np.nan argument as None"""
  51. # GH#17276
  52. s = Series([1, 2, 3])
  53. tm.assert_series_equal(s.clip(np.nan), Series([1, 2, 3]))
  54. tm.assert_series_equal(s.clip(upper=np.nan, lower=np.nan), Series([1, 2, 3]))
  55. # GH#19992
  56. tm.assert_series_equal(s.clip(lower=[0, 4, np.nan]), Series([1, 4, 3]))
  57. tm.assert_series_equal(s.clip(upper=[1, np.nan, 1]), Series([1, 2, 1]))
  58. # GH#40420
  59. s = Series([1, 2, 3])
  60. result = s.clip(0, [np.nan, np.nan, np.nan])
  61. tm.assert_series_equal(s, result)
  62. def test_clip_against_series(self):
  63. # GH#6966
  64. s = Series([1.0, 1.0, 4.0])
  65. lower = Series([1.0, 2.0, 3.0])
  66. upper = Series([1.5, 2.5, 3.5])
  67. tm.assert_series_equal(s.clip(lower, upper), Series([1.0, 2.0, 3.5]))
  68. tm.assert_series_equal(s.clip(1.5, upper), Series([1.5, 1.5, 3.5]))
  69. @pytest.mark.parametrize("inplace", [True, False])
  70. @pytest.mark.parametrize("upper", [[1, 2, 3], np.asarray([1, 2, 3])])
  71. def test_clip_against_list_like(self, inplace, upper):
  72. # GH#15390
  73. original = Series([5, 6, 7])
  74. result = original.clip(upper=upper, inplace=inplace)
  75. expected = Series([1, 2, 3])
  76. if inplace:
  77. result = original
  78. tm.assert_series_equal(result, expected, check_exact=True)
  79. def test_clip_with_datetimes(self):
  80. # GH#11838
  81. # naive and tz-aware datetimes
  82. t = Timestamp("2015-12-01 09:30:30")
  83. s = Series([Timestamp("2015-12-01 09:30:00"), Timestamp("2015-12-01 09:31:00")])
  84. result = s.clip(upper=t)
  85. expected = Series(
  86. [Timestamp("2015-12-01 09:30:00"), Timestamp("2015-12-01 09:30:30")]
  87. )
  88. tm.assert_series_equal(result, expected)
  89. t = Timestamp("2015-12-01 09:30:30", tz="US/Eastern")
  90. s = Series(
  91. [
  92. Timestamp("2015-12-01 09:30:00", tz="US/Eastern"),
  93. Timestamp("2015-12-01 09:31:00", tz="US/Eastern"),
  94. ]
  95. )
  96. result = s.clip(upper=t)
  97. expected = Series(
  98. [
  99. Timestamp("2015-12-01 09:30:00", tz="US/Eastern"),
  100. Timestamp("2015-12-01 09:30:30", tz="US/Eastern"),
  101. ]
  102. )
  103. tm.assert_series_equal(result, expected)
  104. def test_clip_with_timestamps_and_oob_datetimes(self):
  105. # GH-42794
  106. ser = Series([datetime(1, 1, 1), datetime(9999, 9, 9)])
  107. result = ser.clip(lower=Timestamp.min, upper=Timestamp.max)
  108. expected = Series([Timestamp.min, Timestamp.max], dtype="object")
  109. tm.assert_series_equal(result, expected)