undo protocol/other scheduler changes, add asserts for mypy check

Author: ciaconet

src/xdist/dsession.py

@@ -187,13 +187,15 @@ def is_node_finishing(self, node: WorkerController) -> bool:
         Evaluate whether it's on its last test, or if no tests are pending.
         """
         assert self.sched is not None
+        assert type(self.sched) is CustomGroup
         pending = self.sched.node2pending.get(node)
         return pending is not None and len(pending) < 2
 
 
     def is_node_clear(self, node: WorkerController) -> bool:
         """Check if a test worker has no pending tests."""
         assert self.sched is not None
+        assert type(self.sched) is CustomGroup
         pending = self.sched.node2pending.get(node)
         return pending is None or len(pending) == 0
 
@@ -218,13 +220,15 @@ def are_all_active_nodes_collected(self) -> bool:
 
     def reset_nodes_if_needed(self) -> None:
         assert self.sched is not None
+        assert type(self.sched) is CustomGroup
         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) -> None:
         """Issue shutdown notices to workers for rescheduling purposes."""
         assert self.sched is not None
+        assert type(self.sched) is CustomGroup
         if len(self.sched.pending) != 0:
             self.remake_nodes = True
         for node in self.sched.nodes:
@@ -235,12 +239,14 @@ def reset_nodes(self) -> None:
     def reschedule(self) -> None:
         """Reschedule tests."""
         assert self.sched is not None
+        assert type(self.sched) is CustomGroup
         self.sched.do_resched = False
         self.sched.check_schedule(self.sched.nodes[0], 1.0, True)
 
 
     def prepare_for_reschedule(self) -> None:
         """Update test workers and their status tracking so rescheduling is ready."""
+        assert type(self.sched) is CustomGroup
         assert self.sched is not None
         self.remake_nodes = False
         num_workers = self.sched.dist_groups[self.sched.pending_groups[0]]['group_workers']
@@ -251,6 +257,7 @@ def prepare_for_reschedule(self) -> None:
         self._active_nodes = set()
         self._active_nodes.update(new_nodes)
         self.sched.node2pending = {}
+        assert type(self.sched) is CustomGroup
         self.sched.do_resched = True
 
     #

src/xdist/scheduler/each.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from typing import Any
 from typing import Sequence
 
 import pytest
@@ -30,10 +29,6 @@ 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:
@@ -111,9 +106,6 @@ 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,7 +1,6 @@
 from __future__ import annotations
 
 from itertools import cycle
-from typing import Any
 from typing import Sequence
 
 import pytest
@@ -63,9 +62,6 @@ 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")
@@ -181,7 +177,7 @@ def remove_pending_tests_from_node(
         raise NotImplementedError()
 
 
-    def check_schedule(self, node: WorkerController, duration: float = 0, from_dsession: bool = False) -> None:
+    def check_schedule(self, node: WorkerController, duration: float = 0) -> 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,7 +1,6 @@
 from __future__ import annotations
 
 from collections import OrderedDict
-from typing import Any
 from typing import NoReturn
 from typing import Sequence
 
@@ -94,14 +93,10 @@ 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:
@@ -167,9 +162,6 @@ 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,19 +1,12 @@
 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]: ...
 
@@ -34,8 +27,6 @@ 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,6 +1,5 @@
 from __future__ import annotations
 
-from typing import Any
 from typing import NamedTuple
 from typing import Sequence
 
@@ -71,9 +70,6 @@ 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:
@@ -197,8 +193,7 @@ def remove_pending_tests_from_node(
         self.pending.extend(indices)
         self.check_schedule()
 
-    def check_schedule(self, node: WorkerController | None = None, duration: float = 0, from_dsession: bool = False
-                       ) -> None:
+    def check_schedule(self) -> None:
         """Reschedule tests/perform load balancing."""
         nodes_up = [
             NodePending(node, pending)