diff --git a/arrow/__init__.py b/arrow/__init__.py index 9232b379..acd21582 100644 --- a/arrow/__init__.py +++ b/arrow/__init__.py @@ -1,6 +1,7 @@ from ._version import __version__ from .api import get, now, utcnow from .arrow import Arrow +from .duration import Duration from .factory import ArrowFactory from .formatter import ( FORMAT_ATOM, @@ -26,6 +27,7 @@ "utcnow", "Arrow", "ArrowFactory", + "Duration", "FORMAT_ATOM", "FORMAT_COOKIE", "FORMAT_RFC822", diff --git a/arrow/duration.py b/arrow/duration.py new file mode 100644 index 00000000..5eda799b --- /dev/null +++ b/arrow/duration.py @@ -0,0 +1,308 @@ +""" +Provides the :class:`Duration ` class for ISO 8601 +duration parsing, formatting, and arithmetic with Arrow objects. + +ISO 8601 durations follow the format: P[n]Y[n]M[n]DT[n]H[n]M[n]S + +Examples: + P1Y2M3DT4H5M6S -> 1 year, 2 months, 3 days, 4 hours, 5 minutes, 6 seconds + P3D -> 3 days + PT12H -> 12 hours + P1Y6M -> 1 year, 6 months + P0.5Y -> 0.5 years +""" + +import re +from datetime import timedelta +from typing import Optional, Union + +from dateutil.relativedelta import relativedelta + +_ISO8601_DURATION_RE = re.compile( + r"^P" + r"(?:(?P\d+(?:\.\d+)?)Y)?" + r"(?:(?P\d+(?:\.\d+)?)M)?" + r"(?:(?P\d+(?:\.\d+)?)W)?" + r"(?:(?P\d+(?:\.\d+)?)D)?" + r"(?:T" + r"(?:(?P\d+(?:\.\d+)?)H)?" + r"(?:(?P\d+(?:\.\d+)?)M)?" + r"(?:(?P\d+(?:\.\d+)?)S)?" + r")?$" +) + + +class Duration: + """Represents an ISO 8601 duration. + + A duration consists of years, months, weeks, days, hours, minutes, + and seconds. Unlike a :class:`datetime.timedelta`, a Duration can + represent calendar-based intervals (years, months) that don't have + a fixed number of days. + + Usage:: + + >>> from arrow.duration import Duration + >>> d = Duration.parse("P1Y2M3DT4H5M6S") + >>> d.years + 1 + >>> d.months + 2 + >>> d.days + 3 + >>> d.hours + 4 + >>> d.minutes + 5 + >>> d.seconds + 6 + >>> str(d) + 'P1Y2M3DT4H5M6S' + + Durations can be applied to Arrow objects:: + + >>> import arrow + >>> arw = arrow.get("2020-01-01") + >>> d = Duration.parse("P1Y2M") + >>> arw + d + + + """ + + def __init__( + self, + years: Union[int, float] = 0, + months: Union[int, float] = 0, + weeks: Union[int, float] = 0, + days: Union[int, float] = 0, + hours: Union[int, float] = 0, + minutes: Union[int, float] = 0, + seconds: Union[int, float] = 0, + ) -> None: + self.years = years + self.months = months + self.weeks = weeks + self.days = days + self.hours = hours + self.minutes = minutes + self.seconds = seconds + + @classmethod + def parse(cls, iso_string: str) -> "Duration": + """Parse an ISO 8601 duration string into a Duration object. + + Parameters: + iso_string: An ISO 8601 duration string (e.g., ``"P1Y2M3DT4H5M6S"``). + + Returns: + A new :class:`Duration` instance. + + Raises: + ValueError: If the string is not a valid ISO 8601 duration. + """ + if not iso_string: + raise ValueError("Duration string cannot be empty.") + + match = _ISO8601_DURATION_RE.match(iso_string) + if not match: + raise ValueError(f"Invalid ISO 8601 duration: {iso_string!r}") + + groups = match.groupdict() + + # At least one component must be present + if all(v is None for v in groups.values()): + raise ValueError(f"Invalid ISO 8601 duration: {iso_string!r}") + + def _to_num(val: Optional[str]) -> Union[int, float]: + if val is None: + return 0 + f = float(val) + return int(f) if f == int(f) else f + + return cls( + years=_to_num(groups["years"]), + months=_to_num(groups["months"]), + weeks=_to_num(groups["weeks"]), + days=_to_num(groups["days"]), + hours=_to_num(groups["hours"]), + minutes=_to_num(groups["minutes"]), + seconds=_to_num(groups["seconds"]), + ) + + @classmethod + def from_timedelta(cls, td: timedelta) -> "Duration": + """Create a Duration from a :class:`datetime.timedelta`. + + Note: Only days and seconds are preserved. Years and months + cannot be inferred from a timedelta. + """ + total_seconds = int(td.total_seconds()) + days, remainder = divmod(abs(total_seconds), 86400) + hours, remainder = divmod(remainder, 3600) + minutes, seconds = divmod(remainder, 60) + + # Handle sub-second precision + microseconds = td.microseconds + total_seconds_val: Union[int, float] = seconds + if microseconds: + total_seconds_val = seconds + microseconds / 1_000_000 + + return cls( + days=days, + hours=hours, + minutes=minutes, + seconds=total_seconds_val, + ) + + def to_relativedelta(self) -> relativedelta: + """Convert this Duration to a :class:`dateutil.relativedelta.relativedelta`. + + This is used internally for arithmetic with Arrow objects. + """ + return relativedelta( + years=int(self.years), + months=int(self.months), + weeks=int(self.weeks), + days=int(self.days), + hours=int(self.hours), + minutes=int(self.minutes), + seconds=int(self.seconds), + microseconds=int((self.seconds - int(self.seconds)) * 1_000_000), + ) + + def to_timedelta(self) -> timedelta: + """Convert to a :class:`datetime.timedelta`. + + Note: Years are approximated as 365.25 days and months as 30.44 days. + For exact arithmetic, use :meth:`to_relativedelta` or apply the + duration to an Arrow object directly. + """ + total_days = ( + self.years * 365.25 + self.months * 30.44 + self.weeks * 7 + self.days + ) + total_seconds = self.hours * 3600 + self.minutes * 60 + self.seconds + return timedelta(days=total_days, seconds=total_seconds) + + def isoformat(self) -> str: + """Format as an ISO 8601 duration string. + + Returns: + An ISO 8601 duration string (e.g., ``"P1Y2M3DT4H5M6S"``). + """ + date_parts = [] + time_parts = [] + + if self.years: + date_parts.append(f"{_format_num(self.years)}Y") + if self.months: + date_parts.append(f"{_format_num(self.months)}M") + if self.weeks: + date_parts.append(f"{_format_num(self.weeks)}W") + if self.days: + date_parts.append(f"{_format_num(self.days)}D") + + if self.hours: + time_parts.append(f"{_format_num(self.hours)}H") + if self.minutes: + time_parts.append(f"{_format_num(self.minutes)}M") + if self.seconds: + time_parts.append(f"{_format_num(self.seconds)}S") + + result = "P" + "".join(date_parts) + if time_parts: + result += "T" + "".join(time_parts) + + # If everything is zero, return P0D + if result == "P": + result = "P0D" + + return result + + def __str__(self) -> str: + return self.isoformat() + + def __repr__(self) -> str: + parts = [] + for attr in ("years", "months", "weeks", "days", "hours", "minutes", "seconds"): + val = getattr(self, attr) + if val: + parts.append(f"{attr}={val!r}") + return f"Duration({', '.join(parts)})" if parts else "Duration()" + + def __eq__(self, other: object) -> bool: + if isinstance(other, Duration): + return ( + self.years == other.years + and self.months == other.months + and self.weeks == other.weeks + and self.days == other.days + and self.hours == other.hours + and self.minutes == other.minutes + and self.seconds == other.seconds + ) + return NotImplemented + + def __neg__(self) -> "Duration": + return Duration( + years=-self.years, + months=-self.months, + weeks=-self.weeks, + days=-self.days, + hours=-self.hours, + minutes=-self.minutes, + seconds=-self.seconds, + ) + + def __add__(self, other: object) -> Union["Duration", object]: + if isinstance(other, Duration): + return Duration( + years=self.years + other.years, + months=self.months + other.months, + weeks=self.weeks + other.weeks, + days=self.days + other.days, + hours=self.hours + other.hours, + minutes=self.minutes + other.minutes, + seconds=self.seconds + other.seconds, + ) + return NotImplemented + + def __radd__(self, other: object) -> object: + # Support Arrow + Duration + from arrow.arrow import Arrow + + if isinstance(other, Arrow): + return other.shift( + years=int(self.years), + months=int(self.months), + weeks=int(self.weeks), + days=int(self.days), + hours=int(self.hours), + minutes=int(self.minutes), + seconds=int(self.seconds), + ) + return NotImplemented + + def __sub__(self, other: object) -> Union["Duration", object]: + if isinstance(other, Duration): + return self + (-other) + return NotImplemented + + def __hash__(self) -> int: + return hash( + ( + self.years, + self.months, + self.weeks, + self.days, + self.hours, + self.minutes, + self.seconds, + ) + ) + + +def _format_num(val: Union[int, float]) -> str: + """Format a number, dropping unnecessary decimal places.""" + if isinstance(val, float) and val == int(val): + return str(int(val)) + return str(val) diff --git a/tests/test_duration.py b/tests/test_duration.py new file mode 100644 index 00000000..32ede084 --- /dev/null +++ b/tests/test_duration.py @@ -0,0 +1,285 @@ +from datetime import timedelta + +import pytest + +import arrow +from arrow.duration import Duration + + +class TestDurationParse: + def test_full_duration(self): + d = Duration.parse("P1Y2M3DT4H5M6S") + assert d.years == 1 + assert d.months == 2 + assert d.days == 3 + assert d.hours == 4 + assert d.minutes == 5 + assert d.seconds == 6 + + def test_date_only(self): + d = Duration.parse("P1Y2M3D") + assert d.years == 1 + assert d.months == 2 + assert d.days == 3 + assert d.hours == 0 + assert d.minutes == 0 + assert d.seconds == 0 + + def test_time_only(self): + d = Duration.parse("PT4H5M6S") + assert d.years == 0 + assert d.months == 0 + assert d.days == 0 + assert d.hours == 4 + assert d.minutes == 5 + assert d.seconds == 6 + + def test_weeks(self): + d = Duration.parse("P2W") + assert d.weeks == 2 + assert d.days == 0 + + def test_days_only(self): + d = Duration.parse("P30D") + assert d.days == 30 + + def test_hours_only(self): + d = Duration.parse("PT12H") + assert d.hours == 12 + + def test_minutes_only(self): + d = Duration.parse("PT30M") + assert d.minutes == 30 + + def test_seconds_only(self): + d = Duration.parse("PT45S") + assert d.seconds == 45 + + def test_fractional_seconds(self): + d = Duration.parse("PT1.5S") + assert d.seconds == 1.5 + + def test_fractional_hours(self): + d = Duration.parse("PT0.5H") + assert d.hours == 0.5 + + def test_years_months(self): + d = Duration.parse("P1Y6M") + assert d.years == 1 + assert d.months == 6 + + def test_invalid_empty(self): + with pytest.raises(ValueError): + Duration.parse("") + + def test_invalid_no_p(self): + with pytest.raises(ValueError): + Duration.parse("1Y2M3D") + + def test_invalid_only_p(self): + with pytest.raises(ValueError): + Duration.parse("P") + + def test_invalid_only_pt(self): + with pytest.raises(ValueError): + Duration.parse("PT") + + def test_invalid_garbage(self): + with pytest.raises(ValueError): + Duration.parse("not a duration") + + +class TestDurationFormat: + def test_full_format(self): + d = Duration(years=1, months=2, days=3, hours=4, minutes=5, seconds=6) + assert str(d) == "P1Y2M3DT4H5M6S" + + def test_date_only_format(self): + d = Duration(years=1, months=2, days=3) + assert str(d) == "P1Y2M3D" + + def test_time_only_format(self): + d = Duration(hours=4, minutes=5, seconds=6) + assert str(d) == "PT4H5M6S" + + def test_weeks_format(self): + d = Duration(weeks=2) + assert str(d) == "P2W" + + def test_zero_duration(self): + d = Duration() + assert str(d) == "P0D" + + def test_roundtrip(self): + cases = [ + "P1Y2M3DT4H5M6S", + "P3D", + "PT12H", + "P1Y6M", + "P2W", + "PT30M", + "PT45S", + "P1Y", + ] + for case in cases: + assert str(Duration.parse(case)) == case + + def test_isoformat_method(self): + d = Duration(days=5, hours=3) + assert d.isoformat() == "P5DT3H" + + +class TestDurationArithmetic: + def test_add_durations(self): + d1 = Duration(years=1, months=2) + d2 = Duration(years=3, days=5) + result = d1 + d2 + assert result.years == 4 + assert result.months == 2 + assert result.days == 5 + + def test_sub_durations(self): + d1 = Duration(years=3, months=6) + d2 = Duration(years=1, months=2) + result = d1 - d2 + assert result.years == 2 + assert result.months == 4 + + def test_negate(self): + d = Duration(years=1, months=2) + neg = -d + assert neg.years == -1 + assert neg.months == -2 + + def test_add_to_arrow(self): + arw = arrow.get("2020-01-01") + d = Duration(years=1, months=2) + result = arw + d + assert result.year == 2021 + assert result.month == 3 + assert result.day == 1 + + def test_add_days_to_arrow(self): + arw = arrow.get("2020-01-01") + d = Duration(days=31) + result = arw + d + assert result.year == 2020 + assert result.month == 2 + assert result.day == 1 + + def test_add_time_to_arrow(self): + arw = arrow.get("2020-01-01T00:00:00") + d = Duration(hours=25, minutes=30) + result = arw + d + assert result.day == 2 + assert result.hour == 1 + assert result.minute == 30 + + +class TestDurationEquality: + def test_equal(self): + d1 = Duration(years=1, months=2) + d2 = Duration(years=1, months=2) + assert d1 == d2 + + def test_not_equal(self): + d1 = Duration(years=1) + d2 = Duration(years=2) + assert d1 != d2 + + def test_hash(self): + d1 = Duration(years=1, months=2) + d2 = Duration(years=1, months=2) + assert hash(d1) == hash(d2) + s = {d1, d2} + assert len(s) == 1 + + +class TestDurationConversions: + def test_to_timedelta(self): + d = Duration(days=5, hours=3, minutes=30) + td = d.to_timedelta() + assert td == timedelta(days=5, hours=3, minutes=30) + + def test_to_relativedelta(self): + d = Duration(years=1, months=2, days=3) + rd = d.to_relativedelta() + assert rd.years == 1 + assert rd.months == 2 + assert rd.days == 3 + + def test_from_timedelta(self): + td = timedelta(days=5, hours=3, minutes=30, seconds=15) + d = Duration.from_timedelta(td) + assert d.days == 5 + assert d.hours == 3 + assert d.minutes == 30 + assert d.seconds == 15 + assert d.years == 0 + assert d.months == 0 + + +class TestDurationRepr: + def test_repr_full(self): + d = Duration(years=1, months=2, days=3, hours=4, minutes=5, seconds=6) + assert ( + repr(d) + == "Duration(years=1, months=2, days=3, hours=4, minutes=5, seconds=6)" + ) + + def test_repr_empty(self): + d = Duration() + assert repr(d) == "Duration()" + + def test_repr_partial(self): + d = Duration(days=5, hours=3) + assert repr(d) == "Duration(days=5, hours=3)" + + +class TestDurationNotImplemented: + def test_eq_non_duration(self): + d = Duration(years=1) + assert d.__eq__("not a duration") is NotImplemented + + def test_add_non_duration(self): + d = Duration(years=1) + assert d.__add__("not a duration") is NotImplemented + + def test_radd_non_arrow(self): + d = Duration(years=1) + assert d.__radd__("not an arrow") is NotImplemented + + def test_sub_non_duration(self): + d = Duration(years=1) + assert d.__sub__("not a duration") is NotImplemented + + +class TestDurationFromTimedelta: + def test_from_timedelta_with_microseconds(self): + td = timedelta(days=1, seconds=30, microseconds=500000) + d = Duration.from_timedelta(td) + assert d.days == 1 + assert d.seconds == 30.5 + + def test_to_timedelta_with_years_months(self): + d = Duration(years=1, months=6) + td = d.to_timedelta() + expected_days = 1 * 365.25 + 6 * 30.44 + assert td == timedelta(days=expected_days) + + +class TestDurationFormatNum: + def test_format_fractional_value(self): + d = Duration(seconds=1.5) + assert str(d) == "PT1.5S" + + def test_format_whole_float(self): + d = Duration(seconds=2.0) + assert str(d) == "PT2S" + + +class TestDurationImport: + def test_importable_from_arrow(self): + from arrow import Duration as D + + assert D is Duration