Type annotate ParameterSet

Author: bluetech

src/_pytest/compat.py

@@ -1,6 +1,7 @@
 """
 python version compatibility code
 """
+import enum
 import functools
 import inspect
 import os
@@ -33,13 +34,20 @@
 
 if TYPE_CHECKING:
     from typing import Type
+    from typing_extensions import Final
 
 
 _T = TypeVar("_T")
 _S = TypeVar("_S")
 
 
-NOTSET = object()
+# fmt: off
+# Singleton type for NOTSET, as described in:
+# https://www.python.org/dev/peps/pep-0484/#support-for-singleton-types-in-unions
+class NotSetType(enum.Enum):
+    token = 0
+NOTSET = NotSetType.token  # type: Final # noqa: E305
+# fmt: on
 
 MODULE_NOT_FOUND_ERROR = (
     "ModuleNotFoundError" if sys.version_info[:2] >= (3, 6) else "ImportError"

src/_pytest/mark/__init__.py

@@ -1,7 +1,9 @@
 """ generic mechanism for marking and selecting python functions. """
+import typing
 import warnings
 from typing import AbstractSet
 from typing import Optional
+from typing import Union
 
 import attr
 
@@ -31,7 +33,11 @@
 old_mark_config_key = StoreKey[Optional[Config]]()
 
 
-def param(*values, **kw):
+def param(
+    *values: object,
+    marks: "Union[MarkDecorator, typing.Collection[Union[MarkDecorator, Mark]]]" = (),
+    id: Optional[str] = None
+) -> ParameterSet:
     """Specify a parameter in `pytest.mark.parametrize`_ calls or
     :ref:`parametrized fixtures <fixture-parametrize-marks>`.
 
@@ -48,7 +54,7 @@ def test_eval(test_input, expected):
     :keyword marks: a single mark or a list of marks to be applied to this parameter set.
     :keyword str id: the id to attribute to this parameter set.
     """
-    return ParameterSet.param(*values, **kw)
+    return ParameterSet.param(*values, marks=marks, id=id)
 
 
 def pytest_addoption(parser):

src/_pytest/mark/structures.py

@@ -1,11 +1,12 @@
+import collections.abc
 import inspect
+import typing
 import warnings
-from collections import namedtuple
-from collections.abc import MutableMapping
 from typing import Any
 from typing import Iterable
 from typing import List
 from typing import Mapping
+from typing import NamedTuple
 from typing import Optional
 from typing import Sequence
 from typing import Set
@@ -17,20 +18,29 @@
 from .._code import getfslineno
 from ..compat import ascii_escaped
 from ..compat import NOTSET
+from ..compat import NotSetType
+from ..compat import TYPE_CHECKING
+from _pytest.config import Config
 from _pytest.outcomes import fail
 from _pytest.warning_types import PytestUnknownMarkWarning
 
+if TYPE_CHECKING:
+    from _pytest.python import FunctionDefinition
+
+
 EMPTY_PARAMETERSET_OPTION = "empty_parameter_set_mark"
 
 
-def istestfunc(func):
+def istestfunc(func) -> bool:
     return (
         hasattr(func, "__call__")
         and getattr(func, "__name__", "<lambda>") != "<lambda>"
     )
 
 
-def get_empty_parameterset_mark(config, argnames, func):
+def get_empty_parameterset_mark(
+    config: Config, argnames: Sequence[str], func
+) -> "MarkDecorator":
     from ..nodes import Collector
 
     requested_mark = config.getini(EMPTY_PARAMETERSET_OPTION)
@@ -53,16 +63,33 @@ def get_empty_parameterset_mark(config, argnames, func):
         fs,
         lineno,
     )
-    return mark(reason=reason)
-
-
-class ParameterSet(namedtuple("ParameterSet", "values, marks, id")):
+    # Type ignored because MarkDecorator.__call__() is a bit tough to
+    # annotate ATM.
+    return mark(reason=reason)  # type: ignore[no-any-return] # noqa: F723
+
+
+class ParameterSet(
+    NamedTuple(
+        "ParameterSet",
+        [
+            ("values", Sequence[Union[object, NotSetType]]),
+            ("marks", "typing.Collection[Union[MarkDecorator, Mark]]"),
+            ("id", Optional[str]),
+        ],
+    )
+):
     @classmethod
-    def param(cls, *values, marks=(), id=None):
+    def param(
+        cls,
+        *values: object,
+        marks: "Union[MarkDecorator, typing.Collection[Union[MarkDecorator, Mark]]]" = (),
+        id: Optional[str] = None
+    ) -> "ParameterSet":
         if isinstance(marks, MarkDecorator):
             marks = (marks,)
         else:
-            assert isinstance(marks, (tuple, list, set))
+            # TODO(py36): Change to collections.abc.Collection.
+            assert isinstance(marks, (collections.abc.Sequence, set))
 
         if id is not None:
             if not isinstance(id, str):
@@ -73,7 +100,11 @@ def param(cls, *values, marks=(), id=None):
         return cls(values, marks, id)
 
     @classmethod
-    def extract_from(cls, parameterset, force_tuple=False):
+    def extract_from(
+        cls,
+        parameterset: Union["ParameterSet", Sequence[object], object],
+        force_tuple: bool = False,
+    ) -> "ParameterSet":
         """
         :param parameterset:
             a legacy style parameterset that may or may not be a tuple,
@@ -89,10 +120,20 @@ def extract_from(cls, parameterset, force_tuple=False):
         if force_tuple:
             return cls.param(parameterset)
         else:
-            return cls(parameterset, marks=[], id=None)
+            # TODO: Refactor to fix this type-ignore. Currently the following
+            # type-checks but crashes:
+            #
+            #   @pytest.mark.parametrize(('x', 'y'), [1, 2])
+            #   def test_foo(x, y): pass
+            return cls(parameterset, marks=[], id=None)  # type: ignore[arg-type] # noqa: F821
 
     @staticmethod
-    def _parse_parametrize_args(argnames, argvalues, *args, **kwargs):
+    def _parse_parametrize_args(
+        argnames: Union[str, List[str], Tuple[str, ...]],
+        argvalues: Iterable[Union["ParameterSet", Sequence[object], object]],
+        *args,
+        **kwargs
+    ) -> Tuple[Union[List[str], Tuple[str, ...]], bool]:
         if not isinstance(argnames, (tuple, list)):
             argnames = [x.strip() for x in argnames.split(",") if x.strip()]
             force_tuple = len(argnames) == 1
@@ -101,13 +142,23 @@ def _parse_parametrize_args(argnames, argvalues, *args, **kwargs):
         return argnames, force_tuple
 
     @staticmethod
-    def _parse_parametrize_parameters(argvalues, force_tuple):
+    def _parse_parametrize_parameters(
+        argvalues: Iterable[Union["ParameterSet", Sequence[object], object]],
+        force_tuple: bool,
+    ) -> List["ParameterSet"]:
         return [
             ParameterSet.extract_from(x, force_tuple=force_tuple) for x in argvalues
         ]
 
     @classmethod
-    def _for_parametrize(cls, argnames, argvalues, func, config, function_definition):
+    def _for_parametrize(
+        cls,
+        argnames: Union[str, List[str], Tuple[str, ...]],
+        argvalues: Iterable[Union["ParameterSet", Sequence[object], object]],
+        func,
+        config: Config,
+        function_definition: "FunctionDefinition",
+    ) -> Tuple[Union[List[str], Tuple[str, ...]], List["ParameterSet"]]:
         argnames, force_tuple = cls._parse_parametrize_args(argnames, argvalues)
         parameters = cls._parse_parametrize_parameters(argvalues, force_tuple)
         del argvalues
@@ -370,7 +421,7 @@ def __getattr__(self, name: str) -> MarkDecorator:
 MARK_GEN = MarkGenerator()
 
 
-class NodeKeywords(MutableMapping):
+class NodeKeywords(collections.abc.MutableMapping):
     def __init__(self, node):
         self.node = node
         self.parent = node.parent

testing/test_doctest.py

@@ -1051,7 +1051,7 @@ def test_number_precision(self, testdir, config_mode):
             ("1e3", "999"),
             # The current implementation doesn't understand that numbers inside
             # strings shouldn't be treated as numbers:
-            pytest.param("'3.1416'", "'3.14'", marks=pytest.mark.xfail),
+            pytest.param("'3.1416'", "'3.14'", marks=pytest.mark.xfail),  # type: ignore
         ],
     )
     def test_number_non_matches(self, testdir, expression, output):