Type annotate CallSpec2

Author: bluetech

src/_pytest/python.py

@@ -4,13 +4,16 @@
 import inspect
 import os
 import sys
+import typing
 import warnings
 from collections import Counter
 from collections import defaultdict
 from collections.abc import Sequence
 from functools import partial
 from typing import Dict
+from typing import Iterable
 from typing import List
+from typing import Mapping
 from typing import Optional
 from typing import Tuple
 from typing import Union
@@ -35,19 +38,24 @@
 from _pytest.compat import safe_getattr
 from _pytest.compat import safe_isclass
 from _pytest.compat import STRING_TYPES
+from _pytest.compat import TYPE_CHECKING
 from _pytest.config import hookimpl
 from _pytest.deprecated import FUNCARGNAMES
 from _pytest.mark import MARK_GEN
 from _pytest.mark import ParameterSet
 from _pytest.mark.structures import get_unpacked_marks
 from _pytest.mark.structures import Mark
+from _pytest.mark.structures import MarkDecorator
 from _pytest.mark.structures import normalize_mark_list
 from _pytest.outcomes import fail
 from _pytest.outcomes import skip
 from _pytest.pathlib import parts
 from _pytest.warning_types import PytestCollectionWarning
 from _pytest.warning_types import PytestUnhandledCoroutineWarning
 
+if TYPE_CHECKING:
+    from typing_extensions import Literal
+
 
 def pyobj_property(name):
     def get(self):
@@ -784,16 +792,17 @@ def hasnew(obj):
 
 
 class CallSpec2:
-    def __init__(self, metafunc):
+    def __init__(self, metafunc: "Metafunc") -> None:
         self.metafunc = metafunc
-        self.funcargs = {}
-        self._idlist = []
-        self.params = {}
-        self._arg2scopenum = {}  # used for sorting parametrized resources
-        self.marks = []
-        self.indices = {}
-
-    def copy(self):
+        self.funcargs = {}  # type: Dict[str, object]
+        self._idlist = []  # type: List[str]
+        self.params = {}  # type: Dict[str, object]
+        # Used for sorting parametrized resources.
+        self._arg2scopenum = {}  # type: Dict[str, int]
+        self.marks = []  # type: List[Mark]
+        self.indices = {}  # type: Dict[str, int]
+
+    def copy(self) -> "CallSpec2":
         cs = CallSpec2(self.metafunc)
         cs.funcargs.update(self.funcargs)
         cs.params.update(self.params)
@@ -803,25 +812,39 @@ def copy(self):
         cs._idlist = list(self._idlist)
         return cs
 
-    def _checkargnotcontained(self, arg):
+    def _checkargnotcontained(self, arg: str) -> None:
         if arg in self.params or arg in self.funcargs:
             raise ValueError("duplicate {!r}".format(arg))
 
-    def getparam(self, name):
+    def getparam(self, name: str) -> object:
         try:
             return self.params[name]
         except KeyError:
             raise ValueError(name)
 
     @property
-    def id(self):
+    def id(self) -> str:
         return "-".join(map(str, self._idlist))
 
-    def setmulti2(self, valtypes, argnames, valset, id, marks, scopenum, param_index):
+    def setmulti2(
+        self,
+        valtypes: "Mapping[str, Literal['params', 'funcargs']]",
+        argnames: typing.Sequence[str],
+        valset: Iterable[object],
+        id: str,
+        marks: Iterable[Union[Mark, MarkDecorator]],
+        scopenum: int,
+        param_index: int,
+    ) -> None:
         for arg, val in zip(argnames, valset):
             self._checkargnotcontained(arg)
             valtype_for_arg = valtypes[arg]
-            getattr(self, valtype_for_arg)[arg] = val
+            if valtype_for_arg == "params":
+                self.params[arg] = val
+            elif valtype_for_arg == "funcargs":
+                self.funcargs[arg] = val
+            else:
+                assert False, "Unhandled valtype for arg: {}".format(valtype_for_arg)
             self.indices[arg] = param_index
             self._arg2scopenum[arg] = scopenum
         self._idlist.append(id)
@@ -1042,7 +1065,9 @@ def _validate_ids(self, ids, parameters, func_name):
                     )
         return new_ids
 
-    def _resolve_arg_value_types(self, argnames: List[str], indirect) -> Dict[str, str]:
+    def _resolve_arg_value_types(
+        self, argnames: List[str], indirect
+    ) -> Dict[str, "Literal['params', 'funcargs']"]:
         """Resolves if each parametrized argument must be considered a parameter to a fixture or a "funcarg"
         to the function, based on the ``indirect`` parameter of the parametrized() call.
 
@@ -1054,7 +1079,9 @@ def _resolve_arg_value_types(self, argnames: List[str], indirect) -> Dict[str, s
             * "funcargs" if the argname should be a parameter to the parametrized test function.
         """
         if isinstance(indirect, bool):
-            valtypes = dict.fromkeys(argnames, "params" if indirect else "funcargs")
+            valtypes = dict.fromkeys(
+                argnames, "params" if indirect else "funcargs"
+            )  # type: Dict[str, Literal["params", "funcargs"]]
         elif isinstance(indirect, Sequence):
             valtypes = dict.fromkeys(argnames, "funcargs")
             for arg in indirect: