Merge pull request #34 from valohai/prun-typecl

Implement a TypeVar for ParallelRun subclasses
valohai · Oct 11, 2024 · d5af612 · d5af612
2 parents b933089 + 8f4241d
commit d5af612
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/hai/parallel.py b/hai/parallel.py
@@ -34,6 +34,7 @@ def failed_task_names(self) -> Set[str]:
 
 
 RT = TypeVar("RT")
+ParallelRunType = TypeVar("ParallelRunType", bound="BaseParallelRun")
 
 
 class BaseParallelRun:
@@ -86,7 +87,7 @@ def wait(
         self,
         fail_fast: bool = True,
         interval: float = 0.5,
-        callback: Optional[Callable[["BaseParallelRun"], None]] = None,
+        callback: Optional[Callable[[ParallelRunType], None]] = None,
         max_wait: Optional[float] = None,
     ) -> List["ApplyResult[Any]"]:
         """
@@ -128,7 +129,7 @@ def wait(
             had_any_incomplete_task = self._wait_tick(fail_fast)
 
             if callback:
-                callback(self)
+                callback(self)  # type: ignore[arg-type]
 
             # If there were no incomplete tasks left last iteration, quit.
             if not had_any_incomplete_task: