test_arithmetic.py 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. from datetime import timedelta
  2. import numpy as np
  3. import pytest
  4. from pandas import (
  5. Interval,
  6. Timedelta,
  7. Timestamp,
  8. )
  9. @pytest.mark.parametrize("method", ["__add__", "__sub__"])
  10. @pytest.mark.parametrize(
  11. "interval",
  12. [
  13. Interval(Timestamp("2017-01-01 00:00:00"), Timestamp("2018-01-01 00:00:00")),
  14. Interval(Timedelta(days=7), Timedelta(days=14)),
  15. ],
  16. )
  17. @pytest.mark.parametrize(
  18. "delta", [Timedelta(days=7), timedelta(7), np.timedelta64(7, "D")]
  19. )
  20. def test_time_interval_add_subtract_timedelta(interval, delta, method):
  21. # https://github.com/pandas-dev/pandas/issues/32023
  22. result = getattr(interval, method)(delta)
  23. left = getattr(interval.left, method)(delta)
  24. right = getattr(interval.right, method)(delta)
  25. expected = Interval(left, right)
  26. assert result == expected
  27. @pytest.mark.parametrize("interval", [Interval(1, 2), Interval(1.0, 2.0)])
  28. @pytest.mark.parametrize(
  29. "delta", [Timedelta(days=7), timedelta(7), np.timedelta64(7, "D")]
  30. )
  31. def test_numeric_interval_add_timedelta_raises(interval, delta):
  32. # https://github.com/pandas-dev/pandas/issues/32023
  33. msg = "|".join(
  34. [
  35. "unsupported operand",
  36. "cannot use operands",
  37. "Only numeric, Timestamp and Timedelta endpoints are allowed",
  38. ]
  39. )
  40. with pytest.raises((TypeError, ValueError), match=msg):
  41. interval + delta
  42. with pytest.raises((TypeError, ValueError), match=msg):
  43. delta + interval
  44. @pytest.mark.parametrize("klass", [timedelta, np.timedelta64, Timedelta])
  45. def test_timedelta_add_timestamp_interval(klass):
  46. delta = klass(0)
  47. expected = Interval(Timestamp("2020-01-01"), Timestamp("2020-02-01"))
  48. result = delta + expected
  49. assert result == expected
  50. result = expected + delta
  51. assert result == expected