Type annotate main.py and some parts related to collection

Author: bluetech

src/_pytest/config/__init__.py

@@ -806,7 +806,7 @@ def __init__(
             self.cache = None  # type: Optional[Cache]
 
     @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
@@ -109,13 +110,18 @@ def pytest_unconfigure() -> None:
     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):
@@ -362,7 +368,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,
@@ -441,7 +447,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

@@ -4,6 +4,7 @@
 from typing import Optional
 from typing import Union
 
+import py.path
 from pluggy import HookspecMarker
 
 from _pytest.compat import TYPE_CHECKING
@@ -15,7 +16,12 @@
     from _pytest.config.argparsing import Parser
     from _pytest.config import ExitCode
     from _pytest.main import Session
+    from _pytest.nodes import Collector
+    from _pytest.nodes import Item
     from _pytest.python import Metafunc
+    from _pytest.python import Module
+    from _pytest.python import PyCollector
+
 
 hookspec = HookspecMarker("pytest")
 
@@ -228,7 +234,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.
 
@@ -268,7 +274,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
@@ -281,7 +287,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` """