mypy: pytracer.py, the last file in coverage/

4 files changed, 45 insertions, 35 deletions

diff --git a/coverage/collector.py b/coverage/collector.py
index ab743ee3..22471504 100644
--- a/coverage/collector.py
+++ b/coverage/collector.py
@@ -232,7 +232,7 @@ class Collector:
     def reset(self) -> None:
         """Clear collected data, and prepare to collect more."""
         # The trace data we are collecting.
-        self.data: TTraceData = {}      # type: ignore[assignment]
+        self.data: TTraceData = {}
 
         # A dictionary mapping file names to file tracer plugin names that will
         # handle them.
@@ -310,12 +310,12 @@ class Collector:
     #
     # New in 3.12: threading.settrace_all_threads: https://github.com/python/cpython/pull/96681
 
-    def _installation_trace(self, frame: FrameType, event: str, arg: Any) -> TTraceFn:
+    def _installation_trace(self, frame: FrameType, event: str, arg: Any) -> Optional[TTraceFn]:
         """Called on new threads, installs the real tracer."""
         # Remove ourselves as the trace function.
         sys.settrace(None)
         # Install the real tracer.
-        fn = self._start_tracer()
+        fn: Optional[TTraceFn] = self._start_tracer()
         # Invoke the real trace function with the current event, to be sure
         # not to lose an event.
         if fn:
diff --git a/coverage/pytracer.py b/coverage/pytracer.py
index 94d2ecdc..326c50ba 100644
--- a/coverage/pytracer.py
+++ b/coverage/pytracer.py
@@ -8,12 +8,16 @@ from __future__ import annotations
 import atexit
 import dis
 import sys
+import threading
 
-from types import FrameType
-from typing import Any, Callable, Dict, Optional
+from types import FrameType, ModuleType
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, cast
 
 from coverage import env
-from coverage.types import TFileDisposition, TTraceData, TTraceFn, TTracer, TWarnFn
+from coverage.types import (
+    TArc, TFileDisposition, TLineNo, TTraceData, TTraceFileData, TTraceFn,
+    TTracer, TWarnFn,
+)
 
 # We need the YIELD_VALUE opcode below, in a comparison-friendly form.
 RESUME = dis.opmap.get('RESUME')
@@ -59,16 +63,16 @@ class PyTracer(TTracer):
         self.warn: TWarnFn
 
         # The threading module to use, if any.
-        self.threading = None
+        self.threading: Optional[ModuleType] = None
 
-        self.cur_file_data = None
-        self.last_line = 0          # int, but uninitialized.
+        self.cur_file_data: Optional[TTraceFileData] = None
+        self.last_line: TLineNo = 0
         self.cur_file_name: Optional[str] = None
         self.context: Optional[str] = None
         self.started_context = False
 
-        self.data_stack = []
-        self.thread = None
+        self.data_stack: List[Tuple[Optional[TTraceFileData], Optional[str], TLineNo, bool]] = []
+        self.thread: Optional[threading.Thread] = None
         self.stopped = False
         self._activity = False
 
@@ -78,7 +82,7 @@ class PyTracer(TTracer):
 
         # Cache a bound method on the instance, so that we don't have to
         # re-create a bound method object all the time.
-        self._cached_bound_method_trace = self._trace
+        self._cached_bound_method_trace: TTraceFn = self._trace
 
     def __repr__(self) -> str:
         me = id(self)
@@ -109,7 +113,13 @@ class PyTracer(TTracer):
                 f.write(stack)
             f.write("\n")
 
-    def _trace(self, frame: FrameType, event: str, arg_unused: Any) -> Optional[TTraceFn]:
+    def _trace(
+        self,
+        frame: FrameType,
+        event: str,
+        arg: Any,                               # pylint: disable=unused-argument
+        lineno: Optional[TLineNo] = None,       # pylint: disable=unused-argument
+    ) -> Optional[TTraceFn]:
         """The trace function passed to sys.settrace."""
 
         if THIS_FILE in frame.f_code.co_filename:
@@ -164,7 +174,7 @@ class PyTracer(TTracer):
             # Improve tracing performance: when calling a function, both caller
             # and callee are often within the same file. if that's the case, we
             # don't have to re-check whether to trace the corresponding
-            # function (which is a little bit espensive since it involves
+            # function (which is a little bit expensive since it involves
             # dictionary lookups). This optimization is only correct if we
             # didn't start a context.
             filename = frame.f_code.co_filename
@@ -180,7 +190,7 @@ class PyTracer(TTracer):
                     tracename = disp.source_filename
                     assert tracename is not None
                     if tracename not in self.data:
-                        self.data[tracename] = set()
+                        self.data[tracename] = set()    # type: ignore[assignment]
                     self.cur_file_data = self.data[tracename]
                 else:
                     frame.f_trace_lines = False
@@ -206,13 +216,13 @@ class PyTracer(TTracer):
         elif event == 'line':
             # Record an executed line.
             if self.cur_file_data is not None:
-                lineno = frame.f_lineno
+                flineno: TLineNo = frame.f_lineno
 
                 if self.trace_arcs:
-                    self.cur_file_data.add((self.last_line, lineno))
+                    cast(Set[TArc], self.cur_file_data).add((self.last_line, flineno))
                 else:
-                    self.cur_file_data.add(lineno)
-                self.last_line = lineno
+                    cast(Set[TLineNo], self.cur_file_data).add(flineno)
+                self.last_line = flineno
 
         elif event == 'return':
             if self.trace_arcs and self.cur_file_data:
@@ -240,7 +250,7 @@ class PyTracer(TTracer):
                         real_return = True
                 if real_return:
                     first = frame.f_code.co_firstlineno
-                    self.cur_file_data.add((self.last_line, -first))
+                    cast(Set[TArc], self.cur_file_data).add((self.last_line, -first))
 
             # Leaving this function, pop the filename stack.
             self.cur_file_data, self.cur_file_name, self.last_line, self.started_context = (
@@ -248,6 +258,7 @@ class PyTracer(TTracer):
             )
             # Leaving a context?
             if self.started_context:
+                assert self.switch_context is not None
                 self.context = None
                 self.switch_context(None)
         return self._cached_bound_method_trace
@@ -284,12 +295,14 @@ class PyTracer(TTracer):
         # right thread.
         self.stopped = True
 
-        if self.threading and self.thread.ident != self.threading.current_thread().ident:
-            # Called on a different thread than started us: we can't unhook
-            # ourselves, but we've set the flag that we should stop, so we
-            # won't do any more tracing.
-            #self.log("~", "stopping on different threads")
-            return
+        if self.threading:
+            assert self.thread is not None
+            if self.thread.ident != self.threading.current_thread().ident:
+                # Called on a different thread than started us: we can't unhook
+                # ourselves, but we've set the flag that we should stop, so we
+                # won't do any more tracing.
+                #self.log("~", "stopping on different threads")
+                return
 
         if self.warn:
             # PyPy clears the trace function before running atexit functions,
diff --git a/coverage/types.py b/coverage/types.py
index a45b831e..736269e0 100644
--- a/coverage/types.py
+++ b/coverage/types.py
@@ -32,8 +32,8 @@ class TTraceFn(Protocol):
         frame: FrameType,
         event: str,
         arg: Any,
-        lineno: Optional[int] = None  # Our own twist, see collector.py
-    ) -> TTraceFn:
+        lineno: Optional[TLineNo] = None  # Our own twist, see collector.py
+    ) -> Optional[TTraceFn]:
         ...
 
 ## Coverage.py tracing
@@ -63,11 +63,9 @@ class TFileDisposition(Protocol):
 # - If measuring arcs in the C tracer, the values are sets of packed arcs (two
 #   line numbers combined into one integer).
 
-TTraceData = Union[
-    Dict[str, Set[TLineNo]],
-    Dict[str, Set[TArc]],
-    Dict[str, Set[int]],
-]
+TTraceFileData = Union[Set[TLineNo], Set[TArc], Set[int]]
+
+TTraceData = Dict[str, TTraceFileData]
 
 class TTracer(Protocol):
     """Either CTracer or PyTracer."""
diff --git a/tox.ini b/tox.ini
index d14fd535..e42773ab 100644
--- a/tox.ini
+++ b/tox.ini
@@ -101,10 +101,9 @@ setenv =
     C2=coverage/cmdline.py coverage/collector.py coverage/config.py coverage/context.py coverage/control.py
     C3=coverage/data.py coverage/debug.py coverage/disposition.py coverage/env.py coverage/exceptions.py coverage/execfile.py
     C4=coverage/files.py coverage/html.py coverage/inorout.py coverage/jsonreport.py coverage/lcovreport.py coverage/misc.py coverage/multiproc.py coverage/numbits.py
-    C5=coverage/parser.py coverage/phystokens.py coverage/plugin.py coverage/plugin_support.py coverage/python.py
+    C5=coverage/parser.py coverage/phystokens.py coverage/plugin.py coverage/plugin_support.py coverage/python.py coverage/pytracer.py
     C6=coverage/report.py coverage/results.py coverage/sqldata.py coverage/summary.py
     C7=coverage/templite.py coverage/tomlconfig.py coverage/types.py coverage/version.py coverage/xmlreport.py
-    # not done yet: pytracer.py
     TYPEABLE_C={env:C1} {env:C2} {env:C3} {env:C4} {env:C5} {env:C6} {env:C7}
     T1=tests/conftest.py tests/coveragetest.py tests/goldtest.py tests/helpers.py tests/mixins.py tests/osinfo.py
     T2=tests/test_annotate.py tests/test_api.py tests/test_arcs.py tests/test_cmdline.py tests/test_collector.py tests/test_concurrency.py