more typing

Author: mzuenni

bin/constraints.py

@@ -76,7 +76,7 @@ def math_eval(text: str) -> Optional[int | float]:
             # eval is dangerous, but on the other hand we run submission code so this is fine
             text = text.replace("^", "**")
             value = eval(text, {"__builtin__": None})
-            return value if value is isinstance(value, (int, float)) else None
+            return value if isinstance(value, (int, float)) else None
         except (SyntaxError, NameError, TypeError, ZeroDivisionError):
             return None
 

bin/contest.py

@@ -1,14 +1,15 @@
 import config
 
 from pathlib import Path
+from typing import cast, Any, Optional
 
 from util import *
 
 # Read the contest.yaml, if available
-_contest_yaml = None
+_contest_yaml: Optional[dict[str, Any]] = None
 
 
-def contest_yaml():
+def contest_yaml() -> dict[str, Any]:
     global _contest_yaml
     if _contest_yaml is not None:
         return _contest_yaml
@@ -25,22 +26,22 @@ def contest_yaml():
 _problems_yaml = None
 
 
-def problems_yaml():
+def problems_yaml() -> Optional[list[dict[str, Any]]]:
     global _problems_yaml
-    if _problems_yaml:
-        return _problems_yaml
     if _problems_yaml is False:
         return None
+    if _problems_yaml:
+        return _problems_yaml
 
     problemsyaml_path = Path("problems.yaml")
     if not problemsyaml_path.is_file():
         _problems_yaml = False
         return None
     _problems_yaml = read_yaml(problemsyaml_path)
-    return _problems_yaml
+    return cast(list[dict[str, Any]], _problems_yaml)
 
 
-def get_api():
+def get_api() -> str:
     api = config.args.api or contest_yaml().get("api")
     if not api:
         fatal(
@@ -105,7 +106,7 @@ def call_api(method, endpoint, **kwargs):
     return r
 
 
-def call_api_get_json(url):
+def call_api_get_json(url: str):
     r = call_api("GET", url)
     r.raise_for_status()
     try:

bin/fuzz.py

@@ -7,6 +7,7 @@
 import time
 import threading
 from pathlib import Path
+from typing import Any, Optional
 
 import parallel
 from util import *
@@ -24,9 +25,11 @@
 
 
 class GeneratorTask:
-    def __init__(self, fuzz: "Fuzz", t, i, tmp_id):
+    def __init__(self, fuzz: "Fuzz", t: generate.TestcaseRule, i: int, tmp_id: int):
         self.fuzz = fuzz
-        self.generator = t.generator
+        generator = t.generator
+        assert generator is not None
+        self.generator = generator
         self.solution = t.config.solution
         self.i = i
         self.tmp_id = tmp_id
@@ -39,13 +42,13 @@ def __init__(self, fuzz: "Fuzz", t, i, tmp_id):
         self.save_mutex = threading.Lock()
         self.saved = False
 
-    def run(self, bar):
+    def run(self, bar: ProgressBar) -> None:
         if self._run(bar):
             self.fuzz.finish_task(self.tmp_id)
         else:
             self.fuzz.finish_task(self.tmp_id, 1 + len(self.fuzz.submissions))
 
-    def _run(self, bar):
+    def _run(self, bar: ProgressBar) -> bool:
         # GENERATE THE TEST DATA
         dir = Path("fuzz") / f"tmp_id_{str(self.tmp_id)}"
         cwd = self.fuzz.problem.tmpdir / "tool_runs" / dir
@@ -104,7 +107,7 @@ def _run(self, bar):
                 self.fuzz.queue.put(SubmissionTask(self, submission, testcase, self.tmp_id))
         return True
 
-    def save_test(self, bar):
+    def save_test(self, bar: ProgressBar) -> None:
         if self.saved:
             return
         save = False
@@ -122,17 +125,23 @@ def save_test(self, bar):
 
 
 class SubmissionTask:
-    def __init__(self, generator_task, submission, testcase, tmp_id):
+    def __init__(
+        self,
+        generator_task: GeneratorTask,
+        submission: run.Submission,
+        testcase: Testcase,
+        tmp_id: int,
+    ):
         self.generator_task = generator_task
         self.submission = submission
         self.testcase = testcase
         self.tmp_id = tmp_id
 
-    def run(self, bar):
+    def run(self, bar: ProgressBar) -> None:
         self._run(bar)
         self.generator_task.fuzz.finish_task(self.tmp_id)
 
-    def _run(self, bar):
+    def _run(self, bar: ProgressBar) -> None:
         r = run.Run(self.generator_task.fuzz.problem, self.submission, self.testcase)
         localbar = bar.start(f"{self.generator_task.i}: {self.submission.name}")
         result = r.run(localbar)
@@ -155,10 +164,11 @@ def __init__(self, problem: problem.Problem):
         # Filter to only keep valid rules depending on seed without duplicates from count
         added_testcase_rules = set()
 
-        def add_testcase(t):
+        def add_testcase(t: generate.TestcaseRule) -> None:
             if (
                 t.in_is_generated
                 and t.parse_error is None
+                and t.generator is not None
                 and t.generator.uses_seed
                 and t.generator.command_string.strip() not in added_testcase_rules
             ):
@@ -177,7 +187,7 @@ def add_testcase(t):
         # SUBMISSIONS
         self.submissions = self.problem.selected_or_accepted_submissions()
 
-    def run(self):
+    def run(self) -> bool:
         if not has_ryaml:
             error("Fuzzing needs the ruamel.yaml python3 library. Install python[3]-ruamel.yaml.")
             return False
@@ -192,7 +202,7 @@ def run(self):
 
         message("Press CTRL+C to stop\n", "Fuzz", color_type=MessageType.LOG)
 
-        def runner(task: GeneratorTask):
+        def runner(task: GeneratorTask | SubmissionTask) -> None:
             task.run(bar)
 
         # config.args.no_bar = True
@@ -203,7 +213,7 @@ def runner(task: GeneratorTask):
         self.tasks = 0
         self.queue = parallel.new_queue(runner, pin=True)
 
-        def soft_exit(sig, frame):
+        def soft_exit(sig: Any, frame: Any) -> None:
             if self.queue.aborted:
                 fatal("Running interrupted", force=True)
             else:
@@ -240,7 +250,7 @@ def soft_exit(sig, frame):
 
     # finish task from generator with tmp_id
     # also add new tasks if queue becomes too empty
-    def finish_task(self, tmp_id=None, count=1):
+    def finish_task(self, tmp_id: Optional[int] = None, count: int = 1) -> None:
         with self.queue:
             # return tmp_id (and reuse it if all submissions are finished)
             if tmp_id is not None:
@@ -259,18 +269,18 @@ def finish_task(self, tmp_id=None, count=1):
                 self.iteration += 1
                 # 1 new generator tasks which will also create one task per submission
                 new_tasks = 1 + len(self.submissions)
-                tmp_id = min(self.free_tmp_id)
-                self.free_tmp_id.remove(tmp_id)
-                self.tmp_id_count[tmp_id] = new_tasks
+                new_tmp_id = min(self.free_tmp_id)
+                self.free_tmp_id.remove(new_tmp_id)
+                self.tmp_id_count[new_tmp_id] = new_tasks
                 self.tasks += new_tasks
                 self.queue.put(
-                    GeneratorTask(self, testcase_rule, self.iteration, tmp_id),
+                    GeneratorTask(self, testcase_rule, self.iteration, new_tmp_id),
                     priority=1,
                 )
 
     # Write new rule to yaml
     # lock between read and write to ensure that no rule gets lost
-    def save_test(self, command):
+    def save_test(self, command: str) -> None:
         with self.generators_yaml_mutex:
             generators_yaml = self.problem.path / "generators/generators.yaml"
             data = None

bin/parallel.py

@@ -19,7 +19,7 @@ def __init__(self, task: T, priority: int, index: int):
         self.index = index
 
     # Note: heapq uses a min heap, so higher priorities are 'smaller'.
-    def __lt__(self, other):
+    def __lt__(self, other: "QueueItem[T]") -> bool:
         if self.priority != other.priority:
             # python priority queue is a min heap but larger priority
             # items should come first => reverse compare
@@ -45,24 +45,24 @@ def __init__(self, f: Callable[[T], Any], pin: bool):
         # mutex to lock parallel access
         self.mutex = threading.RLock()
 
-    def __enter__(self):
+    def __enter__(self) -> None:
         self.mutex.__enter__()
 
-    def __exit__(self, *args):
+    def __exit__(self, *args: Any) -> None:
         self.mutex.__exit__(*args)
 
     # Add one task. Higher priority => done first
-    def put(self, task: T, priority=0):
+    def put(self, task: T, priority: int = 0) -> None:
         raise Exception("Abstract method")
 
     # By default, do nothing on .join(). This is overridden in ParallelQueue.
-    def join(self):
+    def join(self) -> None:
         return
 
-    def done(self):
+    def done(self) -> None:
         raise Exception("Abstract method")
 
-    def abort(self):
+    def abort(self) -> None:
         self.aborted = True
 
 
@@ -71,7 +71,7 @@ def __init__(self, f: Callable[[T], Any], pin: bool):
         super().__init__(f, pin)
 
     # Add one task. Higher priority => done first
-    def put(self, task: T, priority: int = 0):
+    def put(self, task: T, priority: int = 0) -> None:
         # no task will be handled after self.abort() so skip adding
         if self.aborted:
             return
@@ -80,7 +80,7 @@ def put(self, task: T, priority: int = 0):
         heapq.heappush(self.tasks, QueueItem(task, priority, self.total_tasks))
 
     # Execute all tasks.
-    def done(self):
+    def done(self) -> None:
         if self.pin:
             cores = list(os.sched_getaffinity(0))
             os.sched_setaffinity(0, {cores[0]})
@@ -127,7 +127,7 @@ def __init__(self, f: Callable[[T], Any], pin: bool, num_threads: int):
 
         signal.signal(signal.SIGINT, self._interrupt_handler)
 
-    def _worker(self, cores: Literal[False] | list[int] = False):
+    def _worker(self, cores: Literal[False] | list[int] = False) -> None:
         if cores is not False:
             os.sched_setaffinity(0, cores)
         while True:
@@ -164,10 +164,10 @@ def _worker(self, cores: Literal[False] | list[int] = False):
                 if self.missing == 0:
                     self.all_done.notify_all()
 
-    def _interrupt_handler(self, sig, frame):
+    def _interrupt_handler(self, sig: Any, frame: Any) -> None:
         util.fatal("Running interrupted", force=True)
 
-    def _handle_first_error(self):
+    def _handle_first_error(self) -> None:
         if self.first_error is not None:
             first_error = self.first_error
             self.first_error = None
@@ -177,7 +177,7 @@ def _handle_first_error(self):
             raise first_error
 
     # Add one task. Higher priority => done first
-    def put(self, task: T, priority: int = 0):
+    def put(self, task: T, priority: int = 0) -> None:
         with self.mutex:
             # no task should be added after .done() was called
             assert not self.finish
@@ -189,14 +189,14 @@ def put(self, task: T, priority: int = 0):
                 heapq.heappush(self.tasks, QueueItem(task, priority, self.total_tasks))
                 self.todo.notify()
 
-    def join(self):
+    def join(self) -> None:
         # wait for all current task to be completed
         with self.all_done:
             self.all_done.wait_for(lambda: self.missing == 0)
             self._handle_first_error()
 
     # Wait for all tasks to be done and stop all threads
-    def done(self):
+    def done(self) -> None:
         self.finish = True
 
         # notify all workers with permission to leave main loop
@@ -213,7 +213,7 @@ def done(self):
 
     # Discard all remaining work in the queue and stop all workers.
     # Call done() to join the threads.
-    def abort(self):
+    def abort(self) -> None:
         super().abort()
 
         with self.mutex:
@@ -227,7 +227,7 @@ def abort(self):
                 self.all_done.notify_all()
 
 
-def new_queue(f: Callable[[T], Any], pin: bool = False):
+def new_queue(f: Callable[[T], Any], pin: bool = False) -> AbstractQueue[T]:
     """
     f(task): the function to run on each queue item.
 
@@ -242,7 +242,7 @@ def new_queue(f: Callable[[T], Any], pin: bool = False):
         return SequentialQueue(f, pin)
 
 
-def run_tasks(f: Callable[[T], Any], tasks: Sequence[T], pin: bool = False):
+def run_tasks(f: Callable[[T], Any], tasks: Sequence[T], pin: bool = False) -> None:
     queue = new_queue(f, pin)
     for task in tasks:
         queue.put(task)

bin/problem.py

@@ -769,7 +769,7 @@ def download_samples(p) -> list[tuple[Path, Path]]:
         return [t for t in testcases if isinstance(t, tuple)]
 
     # Returns the list of submissions passed as command-line arguments, or the list of accepted submissions by default.
-    def selected_or_accepted_submissions(problem) -> list["run.Submission"]:
+    def selected_or_accepted_submissions(problem) -> list[run.Submission]:
         submissions = problem.submissions()
         if not submissions:
             return []
@@ -778,7 +778,7 @@ def selected_or_accepted_submissions(problem) -> list["run.Submission"]:
         else:
             return [s for s in submissions if s.expected_verdicts == [verdicts.Verdict.ACCEPTED]]
 
-    def submissions(problem) -> list["run.Submission"] | Literal[False]:
+    def submissions(problem) -> list[run.Submission] | Literal[False]:
         if problem._submissions is not None:
             if problem._submissions is False:
                 return False