typing: pytest_collection

Author: blueyed

src/_pytest/assertion/__init__.py

@@ -7,6 +7,7 @@
 from _pytest.assertion import rewrite
 from _pytest.assertion import truncate
 from _pytest.assertion import util
+from _pytest.main import Session
 
 
 def pytest_addoption(parser):
@@ -91,7 +92,7 @@ def undo():
     return hook
 
 
-def pytest_collection(session):
+def pytest_collection(session: Session) -> None:
     # this hook is only called when test modules are collected
     # so for example not in the master process of pytest-xdist
     # (which does not collect test modules)

src/_pytest/hookspec.py

@@ -1,6 +1,14 @@
 """ hook specifications for pytest plugins, invoked from main.py and builtin plugins.  """
+from typing import Any
+from typing import Optional
+
 from pluggy import HookspecMarker
 
+from _pytest.compat import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from _pytest.main import Session
+
 
 hookspec = HookspecMarker("pytest")
 
@@ -158,7 +166,7 @@ def pytest_load_initial_conftests(early_config, parser, args):
 
 
 @hookspec(firstresult=True)
-def pytest_collection(session):
+def pytest_collection(session: "Session") -> Optional[Any]:
     """Perform the collection protocol for the given session.
 
     Stops at first non-None result, see :ref:`firstresult`.

src/_pytest/logging.py

@@ -5,6 +5,7 @@
 from io import StringIO
 from typing import AbstractSet
 from typing import Dict
+from typing import Generator
 from typing import List
 from typing import Mapping
 
@@ -591,7 +592,7 @@ def _log_cli_enabled(self):
         ) is not None or self._config.getini("log_cli")
 
     @pytest.hookimpl(hookwrapper=True, tryfirst=True)
-    def pytest_collection(self):
+    def pytest_collection(self) -> Generator[None, None, None]:
         with self.live_logs_context():
             if self.log_cli_handler:
                 self.log_cli_handler.set_when("collection")

src/_pytest/main.py

@@ -8,6 +8,8 @@
 from typing import Dict
 from typing import FrozenSet
 from typing import List
+from typing import Optional
+from typing import Union
 
 import attr
 import py
@@ -27,6 +29,8 @@
 
 
 if TYPE_CHECKING:
+    from typing_extensions import Literal  # noqa: F401
+
     from _pytest.python import Package
 
 
@@ -249,7 +253,11 @@ def pytest_cmdline_main(config):
     return wrap_session(config, _main)
 
 
-def _main(config, session):
+def _main(
+    config: Config, session: "Session"
+) -> Optional[
+    Union["Literal"[ExitCode.TESTS_FAILED], "Literal"[ExitCode.NO_TESTS_COLLECTED]]
+]:
     """ default command line protocol for initialization, session,
     running tests and reporting. """
     config.hook.pytest_collection(session=session)
@@ -259,6 +267,7 @@ def _main(config, session):
         return ExitCode.TESTS_FAILED
     elif session.testscollected == 0:
         return ExitCode.NO_TESTS_COLLECTED
+    return None
 
 
 def pytest_collection(session):

src/_pytest/terminal.py

@@ -514,7 +514,7 @@ def _width_of_current_line(self):
             # py < 1.6.0
             return self._tw.chars_on_current_line
 
-    def pytest_collection(self):
+    def pytest_collection(self) -> None:
         if self.isatty:
             if self.config.option.verbose >= 0:
                 self.write("collecting ... ", bold=True)

src/_pytest/warnings.py

@@ -1,8 +1,10 @@
 import sys
 import warnings
 from contextlib import contextmanager
+from typing import Generator
 
 import pytest
+from _pytest.main import Session
 
 
 def _setoption(wmod, arg):
@@ -117,7 +119,7 @@ def pytest_runtest_protocol(item):
 
 
 @pytest.hookimpl(hookwrapper=True, tryfirst=True)
-def pytest_collection(session):
+def pytest_collection(session: Session) -> Generator[None, None, None]:
     config = session.config
     with catch_warnings_for_item(
         config=config, ihook=config.hook, when="collect", item=None