implement ItemizedDateDelta.total()

Author: ariebovenberg

pysrc/whenever/__init__.pyi

@@ -735,6 +735,13 @@ class ItemizedDateDelta:
         ] = "trunc",
         round_increment: int = 1,
     ) -> ItemizedDateDelta: ...
+    def total(
+        self,
+        unit: Literal["years", "months", "weeks", "days"],
+        /,
+        *,
+        relative_to: Date,
+    ) -> int: ...
     def __iter__(self) -> Iterable[int]: ...
     def __len__(self) -> int: ...
     def __getitem__(

pysrc/whenever/_math.py

@@ -48,7 +48,7 @@ def years_diff(_a: _date, b: InterimDate, increment: int, /) -> _AbsoluteDiff:
     assert isinstance(b, _date)
     diff = (_a.year - b.year) // increment * increment
     shift = _replace_year(b, b.year + diff)
-    sign = 1 if diff >= 0 else -1
+    sign = 1 if _a >= b else -1
 
     # Check if we overshot
     if (diff > 0 and resolve_leap_day(shift) > _a) or (
@@ -76,7 +76,7 @@ def months_diff(a: _date, b: InterimDate, increment: int, /) -> _AbsoluteDiff:
         ((a.year - b.year) * 12 + (a.month - b.month)) // increment
     ) * increment
     shift = _add_months(b, diff)
-    sign = 1 if diff >= 0 else -1
+    sign = 1 if a >= resolve_leap_day(b) else -1
 
     # Check if we overshot
     if (diff > 0 and shift > a) or (diff < 0 and shift < a):

pysrc/whenever/_pywhenever.py

@@ -188,11 +188,6 @@ def _calendar_unit_index(u: str) -> int:
         )
 
 
-# if SPHINXBUILD:
-#     type(_CalendarUnitPlural).__repr__= lambda self: "Literal[...]"
-#     type(_CalendarUnitPlural).__str__= lambda self: "Literal[...]"
-#     breakpoint()
-
 # Metaclass ugh...it proved the most lightweight way to achieve the following:
 # Allowing the constructors of many classes to take an ISO string as well as the
 # regular arguments (i.e. the __init__ signature).
@@ -665,6 +660,9 @@ def since(
             trunc_date = resolve_leap_day(trunc)
             limit = (expand_date - trunc_date).days
             remainder = (self._py_date - trunc_date).days
+            assert (limit >= 0 and remainder >= 0) or (
+                limit <= 0 and remainder <= 0
+            )  # they should have the same sign, or be zero
             # DROP-PY39: match-case
             if round_mode == "expand":
                 do_expand = remainder * sign > 0
@@ -1601,6 +1599,7 @@ class TimeDelta(_Base):
     __slots__ = ("_total_ns",)
 
     # FUTURE: allow weeks and days (with 24 hour warning)?
+    # TODO: allow passing py timedelta (also for other classes)
     def __init__(
         self,
         *,
@@ -1650,6 +1649,7 @@ def total(
             "microseconds",
             "nanoseconds",
         ],
+        # TODO: allow other local time types?
         relative_to: ZonedDateTime = _UNSET,
     ) -> float:
         """The total size in the given unit, as a floating point number
@@ -4074,8 +4074,25 @@ def subtract(
     def total(
         self, unit: _CalendarUnitPlural, /, *, relative_to: Date
     ) -> float:
-        """Return the total duration expressed in the specified unit as a float"""
-        raise NotImplementedError()  # TODO
+        """Return the total duration expressed in the specified unit as a float
+
+        >>> ItemizedDateDelta(years=1, months=6).total("months", relative_to=Date(2020, 1, 31))
+        18.0
+        >>> ItemizedDateDelta(days=1000).total("years", relative_to=Date(2020, 4, 10))
+        2.73972602739726
+        """
+        shifted = relative_to.add(self)
+        trunc_amount, trunc_date_interim, expand_date_interim = DIFF_FUNCS[
+            unit
+        ](shifted._py_date, relative_to._py_date, 1)
+
+        trunc_date = resolve_leap_day(trunc_date_interim)
+        expand_date = resolve_leap_day(expand_date_interim)
+
+        return (
+            trunc_amount
+            + ((shifted._py_date - trunc_date) / (expand_date - trunc_date))
+        ) * self._sign
 
     # A private constructor. Checks bounds but *not* signs or presence of > 0 fields.
     @classmethod

tests/test_date.py

@@ -959,6 +959,21 @@ class TestSinceAndUntil:
                 8,
                 {"round_mode": "half_ceil", "round_increment": 8},
             ),
+            # regression tests for negative deltas close to 0
+            (
+                Date(2020, 1, 1),
+                Date(2020, 12, 31),
+                "years",
+                -1,
+                {"round_mode": "expand"},
+            ),
+            (
+                Date(2020, 12, 1),
+                Date(2020, 12, 31),
+                "months",
+                -1,
+                {"round_mode": "expand"},
+            ),
         ],
     )
     def test_single_unit(self, d1: Date, d2: Date, unit, delta: int, kwargs):

tests/test_itemized_date_delta.py

@@ -469,6 +469,49 @@ def test_floor_round_mode_behaves_correctly_on_negative(self):
         ).exact_eq(ItemizedDateDelta(years=-3, months=-10))
 
 
+class TestTotal:
+
+    @pytest.mark.parametrize(
+        "d, relative_to, unit, expected",
+        [
+            (
+                ItemizedDateDelta(years=2, months=3, weeks=4, days=5),
+                Date("2021-12-31"),
+                "months",
+                28.096774193548388,
+            ),
+            (
+                ItemizedDateDelta(weeks=2, days=16),
+                Date("2021-04-30"),
+                "months",
+                1.0,
+            ),
+            (
+                ItemizedDateDelta(weeks=-2, days=-18),
+                Date("2021-04-30"),
+                "years",
+                -0.08767123287671233,
+            ),
+            (
+                ItemizedDateDelta(weeks=-2, days=-18),
+                Date("2021-04-30"),
+                "days",
+                -32,
+            ),
+        ],
+    )
+    def test_valid(
+        self,
+        d: ItemizedDateDelta,
+        relative_to: Date,
+        unit: Literal["years", "months", "weeks", "days"],
+        expected: float,
+    ):
+        assert d.total(unit, relative_to=relative_to) == pytest.approx(
+            expected
+        )
+
+
 def test_abs():
     d = ItemizedDateDelta(days=-5, weeks=-3)
     assert abs(d).exact_eq(ItemizedDateDelta(days=5, weeks=3))