Add types to wptrunner.testloader/testrunner

Author: gsnedders

tools/wptrunner/wptrunner/browsers/base.py

@@ -103,7 +103,7 @@ class Browser:
 
     init_timeout: float = 30
 
-    def __init__(self, logger: StructuredLogger, *, manager_number: int, **kwargs: Any):
+    def __init__(self, logger: StructuredLogger, *, manager_number: int, **kwargs: Any) -> None:
         if kwargs:
             logger.warning(f"Browser.__init__ kwargs: {kwargs!r}")
         super().__init__()
@@ -160,7 +160,7 @@ def executor_browser(self) -> Tuple[Type['ExecutorBrowser'], Mapping[str, Any]]:
         with which it should be instantiated"""
         return ExecutorBrowser, {}
 
-    def check_crash(self, process: int, test: str) -> bool:
+    def check_crash(self, process: Optional[int], test: Optional[str]) -> bool:
         """Check if a crash occured and output any useful information to the
         log. Returns a boolean indicating whether a crash occured."""
         return False
@@ -197,7 +197,7 @@ class ExecutorBrowser:
     but in some cases it may have more elaborate methods for setting
     up the browser from the runner process.
     """
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         for k, v in kwargs.items():
             setattr(self, k, v)
 
@@ -246,7 +246,7 @@ class OutputHandler:
     but sometimes use a wrapper e.g. mozrunner.
     """
 
-    def __init__(self, logger: StructuredLogger, command: List[str], **kwargs: Any):
+    def __init__(self, logger: StructuredLogger, command: List[str], **kwargs: Any) -> None:
         self.logger = logger
         self.command = command
         self.pid: Optional[int] = None
@@ -305,7 +305,7 @@ def __init__(self,
                  base_path: str = "/",
                  env: Optional[Mapping[str, str]] = None,
                  supports_pac: bool = True,
-                 **kwargs: Any):
+                 **kwargs: Any) -> None:
         super().__init__(logger, **kwargs)
 
         if webdriver_binary is None:

tools/wptrunner/wptrunner/expected.py

@@ -1,9 +1,7 @@
-# mypy: allow-untyped-defs
+import os.path
 
-import os
 
-
-def expected_path(metadata_path, test_path):
+def expected_path(metadata_path: str, test_path: str) -> str:
     """Path to the expectation data file for a given test path.
 
     This is defined as metadata_path + relative_test_path + .ini

tools/wptrunner/wptrunner/instruments.py

@@ -1,10 +1,3 @@
-# mypy: allow-untyped-defs
-
-import time
-import threading
-
-from . import mpcontext
-
 """Instrumentation for measuring high-level time spent on various tasks inside the runner.
 
 This is lower fidelity than an actual profile, but allows custom data to be considered,
@@ -26,8 +19,67 @@
     do_teardown()
 """
 
-class NullInstrument:
-    def set(self, stack):
+from __future__ import annotations
+
+import threading
+import time
+from abc import ABCMeta, abstractmethod
+from typing import TYPE_CHECKING, Iterable, Sequence
+
+from . import mpcontext
+
+if TYPE_CHECKING:
+    import multiprocessing
+    import sys
+    from multiprocessing.process import BaseProcess
+    from types import TracebackType
+
+    if sys.version_info >= (3, 10):
+        from typing import TypeAlias
+    else:
+        from typing_extensions import TypeAlias
+
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
+
+
+class AbstractInstrument(metaclass=ABCMeta):
+    @abstractmethod
+    def __enter__(self) -> AbstractInstrumentHandler:
+        ...
+
+    @abstractmethod
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        ...
+
+
+class AbstractInstrumentHandler(metaclass=ABCMeta):
+    @abstractmethod
+    def set(self, stack: Sequence[str]) -> None:
+        """Set the current task to stack
+
+        :param stack: A list of strings defining the current task.
+                      These are interpreted like a stack trace so that ["foo"] and
+                      ["foo", "bar"] both show up as descendants of "foo"
+        """
+        ...
+
+    @abstractmethod
+    def pause(self) -> None:
+        """Stop recording a task on the current thread. This is useful if the thread
+        is purely waiting on the results of other threads"""
+        ...
+
+
+class NullInstrument(AbstractInstrument, AbstractInstrumentHandler):
+    def set(self, stack: Sequence[str]) -> None:
         """Set the current task to stack
 
         :param stack: A list of strings defining the current task.
@@ -36,37 +88,47 @@ def set(self, stack):
         """
         pass
 
-    def pause(self):
+    def pause(self) -> None:
         """Stop recording a task on the current thread. This is useful if the thread
         is purely waiting on the results of other threads"""
         pass
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         return self
 
-    def __exit__(self, *args, **kwargs):
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
         return
 
 
-class InstrumentWriter:
-    def __init__(self, queue):
+_InstrumentQueue: TypeAlias = "multiprocessing.Queue[tuple[str, int | None, float, Sequence[str] | None]]"
+
+
+class InstrumentWriter(AbstractInstrumentHandler):
+    def __init__(
+        self,
+        queue: _InstrumentQueue,
+    ) -> None:
         self.queue = queue
 
-    def set(self, stack):
-        stack.insert(0, threading.current_thread().name)
+    def set(self, stack: Sequence[str]) -> None:
+        stack = [threading.current_thread().name, *stack]
         stack = self._check_stack(stack)
         self.queue.put(("set", threading.current_thread().ident, time.time(), stack))
 
-    def pause(self):
+    def pause(self) -> None:
         self.queue.put(("pause", threading.current_thread().ident, time.time(), None))
 
-    def _check_stack(self, stack):
-        assert isinstance(stack, (tuple, list))
+    def _check_stack(self, stack: Sequence[str]) -> Sequence[str]:
         return [item.replace(" ", "_") for item in stack]
 
 
-class Instrument:
-    def __init__(self, file_path):
+class Instrument(AbstractInstrument):
+    def __init__(self, file_path: str) -> None:
         """Instrument that collects data from multiple threads and sums the time in each
         thread. The output is in the format required by flamegraph.pl to enable visualisation
         of the time spent in each task.
@@ -75,12 +137,10 @@ def __init__(self, file_path):
                             at the path will be overwritten
         """
         self.path = file_path
-        self.queue = None
-        self.current = None
-        self.start_time = None
-        self.instrument_proc = None
+        self.queue: _InstrumentQueue | None = None
+        self.instrument_proc: BaseProcess | None = None
 
-    def __enter__(self):
+    def __enter__(self) -> InstrumentWriter:
         assert self.instrument_proc is None
         assert self.queue is None
         mp = mpcontext.get_context()
@@ -89,16 +149,24 @@ def __enter__(self):
         self.instrument_proc.start()
         return InstrumentWriter(self.queue)
 
-    def __exit__(self, *args, **kwargs):
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        assert self.instrument_proc is not None
+        assert self.queue is not None
         self.queue.put(("stop", None, time.time(), None))
         self.instrument_proc.join()
         self.instrument_proc = None
         self.queue = None
 
-    def run(self):
+    def run(self) -> None:
+        assert self.queue is not None
         known_commands = {"stop", "pause", "set"}
         with open(self.path, "w") as f:
-            thread_data = {}
+            thread_data: dict[int | None, tuple[Sequence[str], float]] = {}
             while True:
                 command, thread, time_stamp, stack = self.queue.get()
                 assert command in known_commands
@@ -107,15 +175,24 @@ def run(self):
                 # before exiting. Otherwise for either 'set' or 'pause' we only need to dump
                 # information from the current stack (if any) that was recording on the reporting
                 # thread (as that stack is no longer active).
-                items = []
+                items: Iterable[tuple[Sequence[str], float]]
                 if command == "stop":
                     items = thread_data.values()
                 elif thread in thread_data:
-                    items.append(thread_data.pop(thread))
+                    items = [thread_data.pop(thread)]
+                else:
+                    items = []
                 for output_stack, start_time in items:
-                    f.write("%s %d\n" % (";".join(output_stack), int(1000 * (time_stamp - start_time))))
+                    f.write(
+                        "%s %d\n"
+                        % (
+                            ";".join(output_stack),
+                            int(1000 * (time_stamp - start_time)),
+                        )
+                    )
 
                 if command == "set":
+                    assert stack is not None
                     thread_data[thread] = (stack, time_stamp)
                 elif command == "stop":
                     break

tools/wptrunner/wptrunner/manifestexpected.py

@@ -1,22 +1,48 @@
+"""Manifest structure used to store expected results of a test.
+
+Each manifest file is represented by an ExpectedManifest that
+has one or more TestNode children, one per test in the manifest.
+Each TestNode has zero or more SubtestNode children, one for each
+known subtest of the test.
+"""
+
 # mypy: allow-untyped-defs
 
+from __future__ import annotations
+
 from collections import deque
+from typing import Mapping, cast, overload
 
+from . import expected
 from .wptmanifest.backends import static
 from .wptmanifest.backends.base import ManifestItem
 
-from . import expected
 
-"""Manifest structure used to store expected results of a test.
+@overload
+def data_cls_getter(
+    output_node: None, visited_node: object
+) -> type[ExpectedManifest]:
+    ...
 
-Each manifest file is represented by an ExpectedManifest that
-has one or more TestNode children, one per test in the manifest.
-Each TestNode has zero or more SubtestNode children, one for each
-known subtest of the test.
-"""
+
+@overload
+def data_cls_getter(
+    output_node: ExpectedManifest, visited_node: object
+) -> type[TestNode]:
+    ...
+
+
+@overload
+def data_cls_getter(
+    output_node: TestNode, visited_node: object
+) -> type[SubtestNode]:
+    ...
 
 
-def data_cls_getter(output_node, visited_node):
+def data_cls_getter(
+    output_node: ExpectedManifest | TestNode | None,
+    visited_node: object,
+) -> type[ExpectedManifest | TestNode | SubtestNode]:
     # visited_node is intentionally unused
     if output_node is None:
         return ExpectedManifest
@@ -492,7 +518,9 @@ def is_empty(self):
         return True
 
 
-def get_manifest(metadata_root, test_path, run_info):
+def get_manifest(
+    metadata_root: str, test_path: str, run_info: Mapping[str, object]
+) -> ExpectedManifest | None:
     """Get the ExpectedManifest for a particular test path, or None if there is no
     metadata stored for that test path.
 
@@ -504,15 +532,22 @@ def get_manifest(metadata_root, test_path, run_info):
     manifest_path = expected.expected_path(metadata_root, test_path)
     try:
         with open(manifest_path, "rb") as f:
-            return static.compile(f,
-                                  run_info,
-                                  data_cls_getter=data_cls_getter,
-                                  test_path=test_path)
+            return cast(
+                "ExpectedManifest",
+                static.compile(  # type: ignore[no-untyped-call]
+                    f,
+                    run_info,
+                    data_cls_getter=data_cls_getter,
+                    test_path=test_path,
+                ),
+            )
     except OSError:
         return None
 
 
-def get_dir_manifest(path, run_info):
+def get_dir_manifest(
+    path: str, run_info: Mapping[str, object]
+) -> DirectoryManifest | None:
     """Get the ExpectedManifest for a particular test path, or None if there is no
     metadata stored for that test path.
 
@@ -522,8 +557,13 @@ def get_dir_manifest(path, run_info):
     """
     try:
         with open(path, "rb") as f:
-            return static.compile(f,
-                                  run_info,
-                                  data_cls_getter=lambda x,y: DirectoryManifest)
+            return cast(
+                "DirectoryManifest",
+                static.compile(  # type: ignore[no-untyped-call]
+                    f,
+                    run_info,
+                    data_cls_getter=lambda x, y: DirectoryManifest,
+                ),
+            )
     except OSError:
         return None