[email-rotation] give extend_rotation an --ensure-weeks flag

This allows for easy cronjob writing. You can e.g., automatically
`--ensure-weeks=16` every month, and `extend_rotation.py` will handle
making sure that you always have rotations covering up to a quarter out.

Tested by running locally on a few different values of `--ensure-weeks`.

Author: gburgessiv

email-rotation/extend_rotation.py

@@ -4,7 +4,9 @@
 import collections
 import dataclasses
 import datetime
+import itertools
 import logging
+import math
 from pathlib import Path
 from typing import Iterator
 
@@ -77,6 +79,30 @@ def generate_additional_rotations(
         least_recent_assignees.extend(people_on_this_rotation)
 
 
+def calculate_rotations_to_cover(
+    when: datetime.datetime,
+    rotation_length_weeks: int,
+    current_rotations: list[rotations.Rotation],
+) -> int:
+    """Calculates how many additional rotations are needed to cover up to `when`."""
+    if not current_rotations:
+        raise ValueError("No existing rotations to extend from.")
+
+    last_rotation = current_rotations[-1]
+    rotation_length = datetime.timedelta(weeks=rotation_length_weeks)
+    end_of_last_rotation = last_rotation.start_time + rotation_length
+
+    if end_of_last_rotation >= when:
+        return 0
+
+    delta = when - end_of_last_rotation
+    # Add one because there's almost certainly a non-day component to `delta`.
+    days_needed = delta.days + 1
+    weeks_needed = days_needed / 7
+    rotations_needed = int(math.ceil(weeks_needed / rotation_length_weeks))
+    return rotations_needed
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Extend a rotation with additional members."
@@ -107,12 +133,6 @@ def parse_args() -> argparse.Namespace:
         than overwriting the existing file.
         """,
     )
-    parser.add_argument(
-        "--num-rotations",
-        type=int,
-        default=5,
-        help="Number of rotations to add. Default is %(default)d.",
-    )
     parser.add_argument(
         "--people-per-rotation",
         type=int,
@@ -124,6 +144,21 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable debug logging.",
     )
+
+    rotation_group = parser.add_mutually_exclusive_group()
+    rotation_group.add_argument(
+        "--num-rotations",
+        type=int,
+        help="Number of rotations to add, defaults to 5.",
+    )
+    rotation_group.add_argument(
+        "--ensure-weeks",
+        type=int,
+        help="""
+        Ensure the rotation schedule covers at least this many weeks into
+        the future.
+        """,
+    )
     return parser.parse_args()
 
 
@@ -136,24 +171,50 @@ def main() -> None:
     )
 
     dry_run: bool = opts.dry_run
-    num_rotations: int = opts.num_rotations
     people_per_rotation: int = opts.people_per_rotation
     rotation_file_path: Path = opts.rotation_file
     rotation_length_weeks: int = opts.rotation_length_weeks
     rotation_members_file_path: Path = opts.rotation_members_file
 
+    now = datetime.datetime.now(tz=datetime.timezone.utc)
     members_file = rotations.RotationMembersFile.parse_file(rotation_members_file_path)
     current_rotation = rotations.RotationFile.parse_file(rotation_file_path)
 
+    # Determine number of rotations based on flags
+    if opts.num_rotations is not None:
+        num_rotations_to_add: int = opts.num_rotations
+    elif opts.ensure_weeks is not None:
+        ensure_weeks: int = opts.ensure_weeks
+        num_rotations_to_add = calculate_rotations_to_cover(
+            now + datetime.timedelta(weeks=ensure_weeks),
+            rotation_length_weeks,
+            current_rotation.rotations,
+        )
+        if num_rotations_to_add == 0:
+            logging.info(
+                "Current rotations already cover the next %d weeks; no new rotations needed.",
+                ensure_weeks,
+            )
+            return
+        logging.info(
+            "Ensuring %d weeks of coverage with %d rotations (%d weeks per rotation)",
+            ensure_weeks,
+            num_rotations_to_add,
+            rotation_length_weeks,
+        )
+    else:
+        # Default to 5 rotations if neither flag is specified
+        num_rotations_to_add = 5
+
     rotation_generator = generate_additional_rotations(
         current_rotation.rotations,
         members_file.members,
         people_per_rotation,
         rotation_length_weeks,
-        now=datetime.datetime.now(tz=datetime.timezone.utc),
+        now=now,
     )
 
-    extra_rotations = [x for x, _ in zip(rotation_generator, range(num_rotations))]
+    extra_rotations = list(itertools.islice(rotation_generator, num_rotations_to_add))
     new_rotations = current_rotation.rotations + extra_rotations
     new_rotations_file = dataclasses.replace(current_rotation, rotations=new_rotations)
 

email-rotation/extend_rotation_test.py

@@ -6,6 +6,7 @@
 
 from extend_rotation import (
     MIN_TIME_UTC,
+    calculate_rotations_to_cover,
     find_most_recent_service_times,
     generate_additional_rotations,
 )
@@ -63,6 +64,67 @@ def test_member_in_multiple_rotations_gets_latest(self) -> None:
         self.assertEqual(result, expected)
 
 
+class TestCalculateRotationsToCover(unittest.TestCase):
+    def test_no_existing_rotations_raises_error(self) -> None:
+        when = datetime.datetime(2025, 6, 1, tzinfo=datetime.timezone.utc)
+        with self.assertRaises(ValueError):
+            calculate_rotations_to_cover(
+                when, rotation_length_weeks=2, current_rotations=[]
+            )
+
+    def test_target_time_before_last_rotation_end_returns_zero(self) -> None:
+        last_rotation_start = datetime.datetime(
+            2025, 5, 1, tzinfo=datetime.timezone.utc
+        )
+        first_rotation_start = last_rotation_start - datetime.timedelta(weeks=2)
+        current_rotations = [
+            Rotation(start_time=first_rotation_start, members=["bob", "charlie"]),
+            Rotation(start_time=last_rotation_start, members=["alice", "bob"]),
+        ]
+
+        # Target time is May 10th, which is before the rotation ends on the
+        # 15th.
+        when = datetime.datetime(2025, 5, 10, tzinfo=datetime.timezone.utc)
+        result = calculate_rotations_to_cover(
+            when, rotation_length_weeks=2, current_rotations=current_rotations
+        )
+        self.assertEqual(result, 0)
+
+    def test_one_rotation_needed(self) -> None:
+        # Last rotation starts May 1st, with 2-week length it ends May 15th
+        last_rotation_start = datetime.datetime(
+            2025, 5, 1, tzinfo=datetime.timezone.utc
+        )
+        current_rotations = [
+            Rotation(start_time=last_rotation_start, members=["alice", "bob"])
+        ]
+
+        # Target time is May 20th, needs one more 2-week rotation, since the
+        # previous one ends on the 15th.
+        when = datetime.datetime(2025, 5, 20, tzinfo=datetime.timezone.utc)
+        result = calculate_rotations_to_cover(
+            when, rotation_length_weeks=2, current_rotations=current_rotations
+        )
+        self.assertEqual(result, 1)
+
+    def test_multiple_rotations_needed(self) -> None:
+        last_rotation_start = datetime.datetime(
+            2025, 5, 1, tzinfo=datetime.timezone.utc
+        )
+        current_rotations = [
+            Rotation(start_time=last_rotation_start, members=["alice", "bob"])
+        ]
+
+        # Target time is June 10th, needs multiple rotations. May 15 to June 10
+        # is about 26 days, which is about 3.7 weeks. With 2-week rotations, we
+        # need 2 rotations (4 weeks total).
+        when = datetime.datetime(2025, 6, 10, tzinfo=datetime.timezone.utc)
+        result = calculate_rotations_to_cover(
+            when, rotation_length_weeks=2, current_rotations=current_rotations
+        )
+        self.assertEqual(result, 2)
+
+
 class TestGenerateAdditionalRotations(unittest.TestCase):
     def test_generate_no_prior_rotations(self) -> None:
         members = ["alice", "bob", "charlie"]