Type annotate main.py and some parts related to collection

Author: bluetech

src/_pytest/config/__init__.py

@@ -771,7 +771,7 @@ def __init__(self, pluginmanager, *, invocation_params=None) -> None:
         )
 
     @property
-    def invocation_dir(self):
+    def invocation_dir(self) -> py.path.local:
         """Backward compatibility"""
         return py.path.local(str(self.invocation_params.dir))
 

src/_pytest/doctest.py

@@ -7,6 +7,7 @@
 import warnings
 from contextlib import contextmanager
 from typing import Dict
+from typing import Iterable
 from typing import List
 from typing import Optional
 from typing import Sequence
@@ -108,13 +109,18 @@ def pytest_unconfigure():
     RUNNER_CLASS = None
 
 
-def pytest_collect_file(path, parent):
+def pytest_collect_file(
+    path: py.path.local, parent
+) -> Optional[Union["DoctestModule", "DoctestTextfile"]]:
     config = parent.config
     if path.ext == ".py":
         if config.option.doctestmodules and not _is_setup_py(config, path, parent):
-            return DoctestModule.from_parent(parent, fspath=path)
+            mod = DoctestModule.from_parent(parent, fspath=path)  # type: DoctestModule
+            return mod
     elif _is_doctest(config, path, parent):
-        return DoctestTextfile.from_parent(parent, fspath=path)
+        txt = DoctestTextfile.from_parent(parent, fspath=path)  # type: DoctestTextfile
+        return txt
+    return None
 
 
 def _is_setup_py(config, path, parent):
@@ -355,7 +361,7 @@ def _get_continue_on_failure(config):
 class DoctestTextfile(pytest.Module):
     obj = None
 
-    def collect(self):
+    def collect(self) -> Iterable[DoctestItem]:
         import doctest
 
         # inspired by doctest.testfile; ideally we would use it directly,
@@ -434,7 +440,7 @@ def _mock_aware_unwrap(obj, stop=None):
 
 
 class DoctestModule(pytest.Module):
-    def collect(self):
+    def collect(self) -> Iterable[DoctestItem]:
         import doctest
 
         class MockAwareDocTestFinder(doctest.DocTestFinder):

src/_pytest/hookspec.py

@@ -1,6 +1,19 @@
 """ hook specifications for pytest plugins, invoked from main.py and builtin plugins.  """
+from typing import List
+from typing import Optional
+from typing import Union
+
+import py.path
 from pluggy import HookspecMarker
 
+from _pytest.compat import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from _pytest.nodes import Collector
+    from _pytest.nodes import Item
+    from _pytest.python import Module
+    from _pytest.python import PyCollector
+
 
 hookspec = HookspecMarker("pytest")
 
@@ -207,7 +220,7 @@ def pytest_collect_directory(path, parent):
     """
 
 
-def pytest_collect_file(path, parent):
+def pytest_collect_file(path: py.path.local, parent) -> "Optional[Collector]":
     """ return collection Node or None for the given path. Any new node
     needs to have the specified ``parent`` as a parent.
 
@@ -247,7 +260,7 @@ def pytest_make_collect_report(collector):
 
 
 @hookspec(firstresult=True)
-def pytest_pycollect_makemodule(path, parent):
+def pytest_pycollect_makemodule(path: py.path.local, parent) -> "Optional[Module]":
     """ return a Module collector or None for the given path.
     This hook will be called for each matching test module path.
     The pytest_collect_file hook needs to be used if you want to
@@ -260,7 +273,9 @@ def pytest_pycollect_makemodule(path, parent):
 
 
 @hookspec(firstresult=True)
-def pytest_pycollect_makeitem(collector, name, obj):
+def pytest_pycollect_makeitem(
+    collector: "PyCollector", name: str, obj
+) -> "Union[None, Item, Collector, List[Union[Item, Collector]]]":
     """ return custom item/collector for a python object in a module, or None.
 
     Stops at first non-None result, see :ref:`firstresult` """