From d577b2cdf8c2a5b4de7da2c80ebc6dcc3475432b Mon Sep 17 00:00:00 2001
From: Ross <xkjq@ymail.com>
Date: Tue, 14 Jul 2026 22:00:26 +0100
Subject: [PATCH] feat: add time limit parsing function and update time_to_run
 parameter to accept string inputs

---
 gen_cons.py          | 29 ++++++++++++++++++++++++++---
 gen_proc.py          | 30 ++++++++++++++++++++++++++----
 test/test_workers.py | 14 +++++++++++++-
 3 files changed, 65 insertions(+), 8 deletions(-)

diff --git a/gen_cons.py b/gen_cons.py
index 76d45fd..cc15a6a 100644
--- a/gen_cons.py
+++ b/gen_cons.py
@@ -507,11 +507,34 @@ def load_workers():
     
 
 
+def parse_time_limit(t) -> int:
+    """Parses a time limit string (e.g. '10h', '30m', '45s') or an integer to seconds."""
+    if isinstance(t, int):
+        return t
+    if isinstance(t, float):
+        return int(t)
+    t = str(t).strip().lower()
+    if not t:
+        return 0
+    if t.isdigit():
+        return int(t)
+    if t.endswith("h"):
+        return int(float(t[:-1]) * 3600)
+    if t.endswith("m"):
+        return int(float(t[:-1]) * 60)
+    if t.endswith("s"):
+        return int(float(t[:-1]))
+    try:
+        return int(float(t))
+    except ValueError:
+        raise ValueError(f"Invalid time format: {t}. Expected format like '10h', '30m', '45s' or a number of seconds.")
+
+
 @app.command()
 def main(
     suspend: bool = False,
     solve: bool = True,
-    time_to_run: int = 60 * 60,
+    time_to_run: str = "1h",
     ratio: float = 0.1,
     start_date: datetime.datetime = ROTA_START_DATE,
     weeks: int = 22,
@@ -698,8 +721,8 @@ def main(
     # Rota.build_workers()
     # Rota.build_model()
 
-    solver_options = {"ratio": ratio, "seconds": time_to_run, "threads": 10}
-    # solver_options = {"seconds": time_to_run, "threads": 10}
+    seconds_to_run = parse_time_limit(time_to_run)
+    solver_options = {"ratio": ratio, "seconds": seconds_to_run, "threads": 10}
 
     # start_time = time.time()
     Rota.build_and_solve(solver_options, export=True, export_with_timestamp=False, solve=solve, solver="appsi_highs")
diff --git a/gen_proc.py b/gen_proc.py
index efb0cf2..813ee73 100644
--- a/gen_proc.py
+++ b/gen_proc.py
@@ -46,12 +46,34 @@ from rota_generator.workers import (
     OutOfProgramme,
 )
 
-# TODO update time_to_run to accept a string like "10h" or "30m" and convert to seconds
+def parse_time_limit(t) -> int:
+    """Parses a time limit string (e.g. '10h', '30m', '45s') or an integer to seconds."""
+    if isinstance(t, int):
+        return t
+    if isinstance(t, float):
+        return int(t)
+    t = str(t).strip().lower()
+    if not t:
+        return 0
+    if t.isdigit():
+        return int(t)
+    if t.endswith("h"):
+        return int(float(t[:-1]) * 3600)
+    if t.endswith("m"):
+        return int(float(t[:-1]) * 60)
+    if t.endswith("s"):
+        return int(float(t[:-1]))
+    try:
+        return int(float(t))
+    except ValueError:
+        raise ValueError(f"Invalid time format: {t}. Expected format like '10h', '30m', '45s' or a number of seconds.")
+
+
 @app.command()
 def main(
     suspend: bool = False,
     solve: bool = True,
-    time_to_run: int = 60 * 60 * 10,
+    time_to_run: str = "10h",
     ratio: float = 0.1,
     start_date: datetime.datetime = "2026-09-07",
     weeks: int = 26,
@@ -567,8 +589,8 @@ def main(
     # Rota.build_workers()
     # Rota.build_model()
 
-    solver_options = {"ratio": ratio, "seconds": time_to_run, "threads": 10}
-    # solver_options = {"seconds": time_to_run, "threads": 10}
+    seconds_to_run = parse_time_limit(time_to_run)
+    solver_options = {"ratio": ratio, "seconds": seconds_to_run, "threads": 10}
 
     # start_time = time.time()
     Rota.build_and_solve(solver_options, export=True, solve=solve, solver="appsi_highs", export_with_timestamp=True)
diff --git a/test/test_workers.py b/test/test_workers.py
index 932e099..cdb9005 100644
--- a/test/test_workers.py
+++ b/test/test_workers.py
@@ -2118,4 +2118,16 @@ def test_shift_dependent_dates_model_definition():
     assert len(worker.shift_end_dates) == 1
     assert isinstance(worker.shift_end_dates[0], ShiftEndDate)
     assert worker.shift_end_dates[0].shift == "a"
-    assert worker.shift_end_dates[0].end_date == datetime.date(2022, 4, 7)
\ No newline at end of file
+    assert worker.shift_end_dates[0].end_date == datetime.date(2022, 4, 7)
+
+
+def test_parse_time_limit():
+    from gen_proc import parse_time_limit
+    assert parse_time_limit(3600) == 3600
+    assert parse_time_limit("3600") == 3600
+    assert parse_time_limit("10h") == 36000
+    assert parse_time_limit("30m") == 1800
+    assert parse_time_limit("45s") == 45
+    assert parse_time_limit(" 1.5h ") == 5400
+    with pytest.raises(ValueError):
+        parse_time_limit("invalid")
\ No newline at end of file