diff --git a/email-rotation/extend_rotation.py b/email-rotation/extend_rotation.py index 70ec491..f99a653 100755 --- a/email-rotation/extend_rotation.py +++ b/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 + # Use total seconds to account for any non-day components of `delta`. + days_needed = math.ceil(delta.total_seconds() / (24 * 60 * 60)) + 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) diff --git a/email-rotation/extend_rotation_test.py b/email-rotation/extend_rotation_test.py index 1bd9858..65adf30 100755 --- a/email-rotation/extend_rotation_test.py +++ b/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"]