mypy fixes

Author: ciaconet

src/xdist/dsession.py

@@ -7,6 +7,7 @@
 import sys
 import traceback
 from typing import Any
+from typing import Callable
 from typing import Sequence
 import warnings
 
@@ -60,14 +61,14 @@ def __init__(self, config: pytest.Config) -> None:
         self._failed_collection_errors: dict[object, bool] = {}
         self._active_nodes: set[WorkerController] = set()
         self._failed_nodes_count = 0
-        self.saved_put = None
+        self.saved_put: Callable[[tuple[str, dict[str, Any]]], None]
         self.remake_nodes = False
         self.ready_to_run_tests = False
         self._max_worker_restart = get_default_max_worker_restart(self.config)
         # summary message to print at the end of the session
         self._summary_report: str | None = None
         self.terminal = config.pluginmanager.getplugin("terminalreporter")
-        self.worker_status: dict[WorkerController, str] = {}
+        self.worker_status: dict[str, str] = {}
         if self.terminal:
             self.trdist = TerminalDistReporter(config)
             config.pluginmanager.register(self.trdist, "terminaldistreporter")
@@ -180,63 +181,71 @@ def loop_once(self) -> None:
             self.triggershutdown()
 
 
-    def is_node_finishing(self, node: WorkerController):
+    def is_node_finishing(self, node: WorkerController) -> bool:
         """Check if a test worker is considered to be finishing.
 
         Evaluate whether it's on its last test, or if no tests are pending.
         """
+        assert self.sched is not None
         pending = self.sched.node2pending.get(node)
         return pending is not None and len(pending) < 2
 
 
-    def is_node_clear(self, node: WorkerController):
+    def is_node_clear(self, node: WorkerController) -> bool:
         """Check if a test worker has no pending tests."""
+        assert self.sched is not None
         pending = self.sched.node2pending.get(node)
         return pending is None or len(pending) == 0
 
 
-    def are_all_nodes_finishing(self):
+    def are_all_nodes_finishing(self) -> bool:
         """Check if all workers are finishing (See 'is_node_finishing' above)."""
+        assert self.sched is not None
         return all(self.is_node_finishing(node) for node in self.sched.nodes)
 
 
-    def are_all_nodes_done(self):
+    def are_all_nodes_done(self) -> bool:
         """Check if all nodes have reported to finish."""
         return all(s == "finished" for s in self.worker_status.values())
 
 
-    def are_all_active_nodes_collected(self):
+    def are_all_active_nodes_collected(self) -> bool:
         """Check if all nodes have reported collection to be complete."""
         if not all(n.gateway.id in self.worker_status for n in self._active_nodes):
             return False
         return all(self.worker_status[n.gateway.id] == "collected" for n in self._active_nodes)
 
 
-    def reset_nodes_if_needed(self):
+    def reset_nodes_if_needed(self) -> None:
+        assert self.sched is not None
         if self.are_all_nodes_finishing() and self.ready_to_run_tests and not self.sched.do_resched:
             self.reset_nodes()
 
 
-    def reset_nodes(self):
+    def reset_nodes(self) -> None:
         """Issue shutdown notices to workers for rescheduling purposes."""
+        assert self.sched is not None
         if len(self.sched.pending) != 0:
             self.remake_nodes = True
         for node in self.sched.nodes:
             if self.is_node_finishing(node):
                 node.shutdown()
 
 
-    def reschedule(self):
+    def reschedule(self) -> None:
         """Reschedule tests."""
+        assert self.sched is not None
         self.sched.do_resched = False
         self.sched.check_schedule(self.sched.nodes[0], 1.0, True)
 
 
-    def prepare_for_reschedule(self):
+    def prepare_for_reschedule(self) -> None:
         """Update test workers and their status tracking so rescheduling is ready."""
+        assert self.sched is not None
         self.remake_nodes = False
         num_workers = self.sched.dist_groups[self.sched.pending_groups[0]]['group_workers']
         self.trdist._status = {}
+        assert self.nodemanager is not None
         new_nodes = self.nodemanager.setup_nodes(self.saved_put, num_workers)
         self.worker_status = {}
         self._active_nodes = set()
@@ -310,7 +319,7 @@ def worker_workerfinished(self, node: WorkerController) -> None:
                 assert not crashitem, (crashitem, node)
         self._active_nodes.remove(node)
 
-    def update_worker_status(self, node, status):
+    def update_worker_status(self, node: WorkerController, status: str) -> None:
         """Track the worker status.
 
         Can be used at callbacks like 'worker_workerfinished' so we remember wchic event

src/xdist/scheduler/customgroup.py

@@ -205,7 +205,7 @@ def remove_pending_tests_from_node(
     ) -> None:
         raise NotImplementedError()
 
-    def check_schedule(self, node: WorkerController, duration: float = 0, from_dsession=False) -> None:
+    def check_schedule(self, node: WorkerController, duration: float = 0, from_dsession: bool = False) -> None:
         """Maybe schedule new items on the node.
 
         If there are any globally pending nodes left then this will
@@ -228,7 +228,7 @@ def check_schedule(self, node: WorkerController, duration: float = 0, from_dsess
                     dist_group_key = self.pending_groups.pop(0)
                     dist_group = self.dist_groups[dist_group_key]
                     nodes = cycle(self.nodes[0:dist_group['group_workers']])
-                    schedule_log = {n.gateway.id:[] for n in self.nodes[0:dist_group['group_workers']]}
+                    schedule_log: dict[str, Any] = {n.gateway.id:[] for n in self.nodes[0:dist_group['group_workers']]}
                     for _ in range(len(dist_group['test_indices'])):
                         n = next(nodes)
                         #needs cleaner way to be identified
@@ -242,7 +242,7 @@ def check_schedule(self, node: WorkerController, duration: float = 0, from_dsess
                     self.report_line(message)
 
         else:
-            pending = self.node2pending.get(node)
+            pending = self.node2pending.get(node, [])
             if len(pending) < 2:
                 self.report_line(
                     f"[-] [csg] Shutting down {node.workerinput['workerid']} because only one case is pending"
@@ -306,7 +306,7 @@ def schedule(self) -> None:
         if not self.collection:
             return
 
-        dist_groups = {}
+        dist_groups: dict[str, dict[Any, Any]] = {}
 
         if self.is_first_time:
             for i, test in enumerate(self.collection):
@@ -343,7 +343,7 @@ def schedule(self) -> None:
         dist_group_key = self.pending_groups.pop(0)
         dist_group = self.dist_groups[dist_group_key]
         nodes = cycle(self.nodes[0:dist_group['group_workers']])
-        schedule_log = {n.gateway.id: [] for n in self.nodes[0:dist_group['group_workers']]}
+        schedule_log: dict[str, Any] = {n.gateway.id: [] for n in self.nodes[0:dist_group['group_workers']]}
         for _ in range(len(dist_group['test_indices'])):
             n = next(nodes)
             # needs cleaner way to be identified
@@ -362,7 +362,7 @@ def _send_tests(self, node: WorkerController, num: int) -> None:
             self.node2pending[node].extend(tests_per_node)
             node.send_runtest_some(tests_per_node)
 
-    def _send_tests_group(self, node: WorkerController, num: int, dist_group_key) -> None:
+    def _send_tests_group(self, node: WorkerController, num: int, dist_group_key: str) -> None:
         tests_per_node = self.dist_groups[dist_group_key]['pending_indices'][:num]
         if tests_per_node:
             del self.dist_groups[dist_group_key]['pending_indices'][:num]

src/xdist/scheduler/each.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from typing import Any
 from typing import Sequence
 
 import pytest
@@ -29,6 +30,10 @@ def __init__(self, config: pytest.Config, log: Producer | None = None) -> None:
         self.numnodes = len(parse_spec_config(config))
         self.node2collection: dict[WorkerController, list[str]] = {}
         self.node2pending: dict[WorkerController, list[int]] = {}
+        self.do_resched: bool = False
+        self.pending: list[int] = []
+        self.dist_groups: dict[str, Any] = {}
+        self.pending_groups: list[str] = []
         self._started: list[WorkerController] = []
         self._removed2pending: dict[WorkerController, list[int]] = {}
         if log is None:
@@ -106,6 +111,9 @@ def add_node_collection(
                     self.node2pending[node] = pending
                     break
 
+    def check_schedule(self, node: WorkerController, duration: float = 0, from_dsession: bool = False) -> None:
+        raise NotImplementedError()
+
     def mark_test_complete(
         self, node: WorkerController, item_index: int, duration: float = 0
     ) -> None:

src/xdist/scheduler/load.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from itertools import cycle
+from typing import Any
 from typing import Sequence
 
 import pytest
@@ -62,6 +63,9 @@ def __init__(self, config: pytest.Config, log: Producer | None = None) -> None:
         self.node2collection: dict[WorkerController, list[str]] = {}
         self.node2pending: dict[WorkerController, list[int]] = {}
         self.pending: list[int] = []
+        self.do_resched: bool = False
+        self.dist_groups: dict[str, Any] = {}
+        self.pending_groups: list[str] = []
         self.collection: list[str] | None = None
         if log is None:
             self.log = Producer("loadsched")
@@ -176,7 +180,8 @@ def remove_pending_tests_from_node(
     ) -> None:
         raise NotImplementedError()
 
-    def check_schedule(self, node: WorkerController, duration: float = 0) -> None:
+
+    def check_schedule(self, node: WorkerController, duration: float = 0, from_dsession: bool = False) -> None:
         """Maybe schedule new items on the node.
 
         If there are any globally pending nodes left then this will

src/xdist/scheduler/loadscope.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections import OrderedDict
+from typing import Any
 from typing import NoReturn
 from typing import Sequence
 
@@ -93,11 +94,14 @@ class LoadScopeScheduling:
     def __init__(self, config: pytest.Config, log: Producer | None = None) -> None:
         self.numnodes = len(parse_spec_config(config))
         self.collection: list[str] | None = None
-
+        self.node2pending: dict[WorkerController, list[int]] = {}
         self.workqueue: OrderedDict[str, dict[str, bool]] = OrderedDict()
         self.assigned_work: dict[WorkerController, dict[str, dict[str, bool]]] = {}
         self.registered_collections: dict[WorkerController, list[str]] = {}
-
+        self.do_resched: bool = False
+        self.pending: list[int] = []
+        self.dist_groups: dict[str, Any] = {}
+        self.pending_groups: list[str] = []
         if log is None:
             self.log = Producer("loadscopesched")
         else:
@@ -163,6 +167,9 @@ def add_node(self, node: WorkerController) -> None:
         assert node not in self.assigned_work
         self.assigned_work[node] = {}
 
+    def check_schedule(self, node: WorkerController, duration: float = 0, from_dsession: bool = False) -> None:
+        raise NotImplementedError()
+
     def remove_node(self, node: WorkerController) -> str | None:
         """Remove a node from the scheduler.
 

src/xdist/scheduler/protocol.py

@@ -1,12 +1,19 @@
 from __future__ import annotations
 
+from typing import Any
 from typing import Protocol
 from typing import Sequence
 
 from xdist.workermanage import WorkerController
 
 
 class Scheduling(Protocol):
+    node2pending: Any
+    do_resched: bool
+    pending: list[int]
+    dist_groups: dict[str, Any]
+    pending_groups: list[str]
+
     @property
     def nodes(self) -> list[WorkerController]: ...
 
@@ -27,6 +34,8 @@ def add_node_collection(
         collection: Sequence[str],
     ) -> None: ...
 
+    def check_schedule(self, node: WorkerController, duration: float = 0, from_dsession: bool = False) -> None: ...
+
     def mark_test_complete(
         self,
         node: WorkerController,

src/xdist/scheduler/worksteal.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from typing import Any
 from typing import NamedTuple
 from typing import Sequence
 
@@ -70,6 +71,9 @@ def __init__(self, config: pytest.Config, log: Producer | None = None) -> None:
         self.node2pending: dict[WorkerController, list[int]] = {}
         self.pending: list[int] = []
         self.collection: list[str] | None = None
+        self.do_resched: bool = False
+        self.dist_groups: dict[str, Any] = {}
+        self.pending_groups: list[str] = []
         if log is None:
             self.log = Producer("workstealsched")
         else:
@@ -193,7 +197,8 @@ def remove_pending_tests_from_node(
         self.pending.extend(indices)
         self.check_schedule()
 
-    def check_schedule(self) -> None:
+    def check_schedule(self, node: WorkerController | None = None, duration: float = 0, from_dsession: bool = False
+                       ) -> None:
         """Reschedule tests/perform load balancing."""
         nodes_up = [
             NodePending(node, pending)