Merge pull request #23 from George-Ogden/mock-patch

Typed Patching

Author: George-Ogden

.github/workflows/test.yaml

@@ -17,12 +17,12 @@ jobs:
     uses: George-Ogden/actions/.github/workflows/python-test.yaml@v3.1.3
     with:
       python-versions: "['3.13']"
-      timeout-minutes: 10
+      timeout-minutes: 15
       pytest-flags: -vv mypy_pytest_plugin -n auto --shard-id=${{ matrix.shard }} --num-shards=${{ matrix.total-shards }}
 
   integration_tests:
     uses: George-Ogden/actions/.github/workflows/python-test.yaml@v3.1.3
     with:
       python-versions: "['3.13']"
-      timeout-minutes: 5
+      timeout-minutes: 10
       pytest-flags: -vv tests

mypy_pytest_plugin/argmapper.py

@@ -0,0 +1,76 @@
+from collections.abc import Collection
+import functools
+from typing import Final
+
+from mypy.argmap import map_formals_to_actuals
+from mypy.checker import TypeChecker
+from mypy.nodes import ArgKind, CallExpr, Expression
+from mypy.typeops import bind_self
+from mypy.types import CallableType, FunctionLike, Instance, Overloaded, Type
+
+type ArgMap = dict[str, Expression]
+
+
+class ArgMapper:
+    ACCEPTED_ARG_KINDS: Final[Collection[ArgKind]] = (ArgKind.ARG_POS, ArgKind.ARG_NAMED)
+
+    @classmethod
+    def named_arg_mapping(cls, call: CallExpr, checker: TypeChecker) -> ArgMap:
+        callee_type = checker.lookup_type(call.callee)
+        return cls._named_arg_type_mapping(call, callee_type, checker)
+
+    @classmethod
+    def _named_arg_type_mapping(
+        cls, call: CallExpr, callee_type: Type, checker: TypeChecker
+    ) -> ArgMap:
+        if isinstance(callee_type, CallableType):
+            return cls._named_arg_callable_mapping(call, callee_type, checker)
+        if isinstance(callee_type, Overloaded):
+            return cls._named_arg_overloaded_mapping(call, callee_type, checker)
+        if (
+            isinstance(callee_type, Instance)
+            and (call_node := callee_type.type.names.get("__call__")) is not None
+            and call_node.type is not None
+        ):
+            type_ = call_node.type
+            if isinstance(type_, FunctionLike):
+                type_ = bind_self(type_, callee_type)
+                return cls._named_arg_type_mapping(call, type_, checker)
+        return {}
+
+    @classmethod
+    def _named_arg_callable_mapping(
+        cls, call: CallExpr, callee_type: CallableType, checker: TypeChecker
+    ) -> ArgMap:
+        mapping = map_formals_to_actuals(
+            actual_kinds=call.arg_kinds,
+            actual_names=call.arg_names,
+            formal_kinds=callee_type.arg_kinds,
+            formal_names=callee_type.arg_names,
+            actual_arg_type=lambda i: call.args[i].accept(checker.expr_checker),
+        )
+
+        return {
+            arg_name: call.args[actual_idx]
+            for actual_idx, formal_idxs in enumerate(mapping)
+            if len(formal_idxs) == 1
+            and callee_type.arg_kinds[formal_idx := formal_idxs[0]] in cls.ACCEPTED_ARG_KINDS
+            and formal_idx < len(callee_type.arg_names)
+            and (arg_name := callee_type.arg_names[formal_idx]) is not None
+        }
+
+    @classmethod
+    def _named_arg_overloaded_mapping(
+        cls, call: CallExpr, callee_type: Overloaded, checker: TypeChecker
+    ) -> ArgMap:
+        return functools.reduce(
+            cls._merge_mappings,
+            (
+                cls._named_arg_callable_mapping(call, callable_type, checker)
+                for callable_type in callee_type.items
+            ),
+        )
+
+    @classmethod
+    def _merge_mappings(cls, this: ArgMap, that: ArgMap) -> ArgMap:
+        return {key: expr for key, expr in this.items() if that.get(key, None) is expr}

mypy_pytest_plugin/argmapper_test.py

@@ -0,0 +1,185 @@
+from typing import Any, cast
+
+from mypy.nodes import CallExpr, Expression
+
+from .argmapper import ArgMapper
+from .test_utils import dump_expr, parse
+
+
+def _named_arg_mapping_test_body(defs: str, expected_keys: list[str]) -> None:
+    parse_result = parse(defs)
+    parse_result.accept_all()
+
+    call = parse_result.defs["call"]
+    assert isinstance(call, CallExpr)
+
+    raw_arg_map = ArgMapper.named_arg_mapping(call, parse_result.checker)
+
+    def dump_arg_map(
+        arg_map: dict[str, Expression],
+    ) -> dict[str, tuple[type, dict[str, Any]]]:
+        return {key: (dump_expr(expr)) for key, expr in arg_map.items()}
+
+    assert dump_arg_map(raw_arg_map) == dump_arg_map(
+        {key: cast(Expression, parse_result.defs[key]) for key in expected_keys}
+    )
+
+
+def test_named_arg_mapping_no_args() -> None:
+    _named_arg_mapping_test_body(
+        """
+        def main() -> int:
+            return 0
+
+        call = main()
+        """,
+        [],
+    )
+
+
+def test_named_arg_mapping_no_named_args() -> None:
+    _named_arg_mapping_test_body(
+        """
+        def main(x: int, y: str, /, *args: bool) -> int:
+            return 0
+
+        call = main(3, "2", True, False)
+        """,
+        [],
+    )
+
+
+def test_named_arg_mapping_named_args_only() -> None:
+    _named_arg_mapping_test_body(
+        """
+        def main(x: int, y: bool, *, z: str) -> int:
+            return 0
+
+        call = main(0, z="1", y=False)
+        x = 0
+        y = False
+        z = "1"
+        """,
+        ["x", "y", "z"],
+    )
+
+
+def test_named_arg_mapping_arg_mix() -> None:
+    _named_arg_mapping_test_body(
+        """
+        from typing import Any
+
+        def main(a: int, /, b: bool, *args: Any, c: float, **kwargs: dict) -> int:
+            return 0
+
+        call = main(0, True, 2j, c=3.0, d={}, e=dict())
+        b = True
+        c = 3.0
+        """,
+        ["b", "c"],
+    )
+
+
+def test_named_arg_mapping_simple_overload() -> None:
+    _named_arg_mapping_test_body(
+        """
+        from typing import overload
+
+        @overload
+        def foo(x: int) -> int:
+            ...
+
+        @overload
+        def foo(x: str) -> str:
+            ...
+
+        def foo(x: str | int) -> str | int:
+            return x
+
+        call = foo(x="str")
+        x = "str"
+        """,
+        ["x"],
+    )
+
+
+def test_named_arg_mapping_complex_overload() -> None:
+    _named_arg_mapping_test_body(
+        """
+        from typing import overload
+
+        @overload
+        def foo(x: int, y: str, *, z: bool) -> int:
+            ...
+
+        @overload
+        def foo(x: int, y: str) -> int:
+            ...
+
+        def foo(x: int, y: str, **kwargs: bool) -> int:
+            return 0
+
+        call = foo(y="0", z=False, x=2)
+        x = 2
+        y = "0"
+        """,
+        ["x", "y"],
+    )
+
+
+def test_named_arg_mapping_varargs_varkwargs_overload() -> None:
+    _named_arg_mapping_test_body(
+        """
+        from typing import overload
+
+        @overload
+        def foo(x: int, y: str, **kwargs: bool) -> int:
+            ...
+
+        @overload
+        def foo(x: int, y: str, **kwargs: int) -> int:
+            ...
+
+        def foo(x: int, y: str, **kwargs: bool | int) -> int:
+            return 0
+
+        call = foo(*(1, "2"), z=True)
+        """,
+        [],
+    )
+
+
+def test_named_arg_mapping_instance_method() -> None:
+    _named_arg_mapping_test_body(
+        """
+        from typing import overload
+
+        class Foo:
+            def bar(self, x: str, y: str, *args: str) -> str:
+                return x + y
+
+        foo = Foo()
+        call = foo.bar("a", "b", "c")
+        x = "a"
+        y = "b"
+        """,
+        ["x", "y"],
+    )
+
+
+def test_named_arg_mapping_call_method() -> None:
+    _named_arg_mapping_test_body(
+        """
+        from typing import overload
+
+        class Foo:
+            def __call__(self, x: str, y: str, *args: str) -> str:
+                return x + y
+
+        foo = Foo()
+        call = foo("a", "b", "c")
+        x = "a"
+        y = "b"
+        """,
+        ["x", "y"],
+    )

mypy_pytest_plugin/decorator_wrapper.py

@@ -1,13 +1,13 @@
-from collections.abc import Collection, Sequence
+from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import Self, TypeGuard
 
-from mypy.argmap import map_actuals_to_formals
 from mypy.checker import TypeChecker
-from mypy.nodes import ArgKind, CallExpr, Expression
-from mypy.types import CallableType, Instance, Type
+from mypy.nodes import CallExpr, Expression
+from mypy.types import Instance
 
-from .error_codes import VARIADIC_ARGNAMES_ARGVALUES
+from .argmapper import ArgMapper
+from .error_codes import UNREADABLE_ARGNAMES_ARGVALUES
 
 
 @dataclass(frozen=True, slots=True)
@@ -37,54 +37,15 @@ def _is_parametrized_decorator_node(
             )
         return False
 
-    def _get_arg_type(self, i: int) -> Type:
-        # subtract one for self
-        i -= 1
-        return self.call.args[i].accept(self.checker.expr_checker)
-
     @property
     def arg_names_and_arg_values(self) -> tuple[Expression, Expression] | None:
-        mapping = map_actuals_to_formals(
-            actual_kinds=[ArgKind.ARG_POS, *self.call.arg_kinds],
-            actual_names=[None, *self.call.arg_names],
-            formal_kinds=self.fn_type.arg_kinds,
-            formal_names=self.fn_type.arg_names,
-            actual_arg_type=self._get_arg_type,
-        )
-        return self._check_actuals_formals_mapping(mapping)
-
-    @property
-    def fn_type(self) -> CallableType:
-        callee_type = self.call.callee.accept(self.checker.expr_checker)
-        assert isinstance(callee_type, Instance)
-        fn_type = callee_type.type.names["__call__"].type
-        assert isinstance(fn_type, CallableType)
-        return fn_type
-
-    def _check_actuals_formals_mapping(
-        self, mapping: list[list[int]]
-    ) -> tuple[Expression, Expression] | None:
-        arg_names_idx, arg_values_idx, *_ = self._clean_up_actuals_formals_mapping(mapping)
-        if (
-            self.call.arg_kinds[arg_values_idx] in self.accepted_arg_kinds
-            and self.call.arg_kinds[arg_names_idx] in self.accepted_arg_kinds
-        ):
-            return self.call.args[arg_names_idx], self.call.args[arg_values_idx]
-        self.checker.fail(
-            "Unable to read argnames and argvalues in a variadic argument.",
-            context=self.call,
-            code=VARIADIC_ARGNAMES_ARGVALUES,
-        )
-        return None
-
-    def _clean_up_actuals_formals_mapping(
-        self, mapping: list[list[int]]
-    ) -> tuple[int, int, list[list[int]]]:
-        [_, [arg_names_idx], [arg_values_idx], *extras] = mapping
-        arg_values_idx -= 1
-        arg_names_idx -= 1
-        return arg_names_idx, arg_values_idx, extras
-
-    @property
-    def accepted_arg_kinds(self) -> Collection[ArgKind]:
-        return (ArgKind.ARG_POS, ArgKind.ARG_NAMED)
+        name_mapping = ArgMapper.named_arg_mapping(self.call, self.checker)
+        try:
+            return name_mapping["argnames"], name_mapping["argvalues"]
+        except KeyError:
+            self.checker.fail(
+                "Unable to read argnames and argvalues. Use positional or keyword arguments.",
+                context=self.call,
+                code=UNREADABLE_ARGNAMES_ARGVALUES,
+            )
+            return None

mypy_pytest_plugin/error_codes.py

@@ -39,9 +39,9 @@
     "missing-argname", "Argument not included in Pytest parametrization.", category="Pytest"
 )
 
-VARIADIC_ARGNAMES_ARGVALUES: Final[ErrorCode] = ErrorCode(
-    "variadic-argnames-argvalues",
-    "Unable to parse variadic argnames or argvalues.",
+UNREADABLE_ARGNAMES_ARGVALUES: Final[ErrorCode] = ErrorCode(
+    "unreadable-argnames-argvalues",
+    "Unable to read argnames or argvalues.",
     category="Pytest",
 )
 

mypy_pytest_plugin/excluded_test_checker_test.py

@@ -41,8 +41,7 @@ def test_6() -> None: ...
         """
     )
 
-    for statement in parse_result.raw_defs:
-        statement.accept(parse_result.checker)
+    parse_result.accept_all()
     ignored_tests = ExcludedTestChecker.ignored_test_names(
         parse_result.raw_defs, parse_result.checker
     )

mypy_pytest_plugin/fixture_manager_test.py

@@ -94,8 +94,7 @@ def _fixture_manager_resolve_requests_and_fixtures_test_body(
     checker = parse_result.checkers[last_module_name]
     fixture_def = parse_result.defs[fullname]
     assert isinstance(fixture_def, FuncDef)
-    for def_ in parse_result.raw_defs:
-        def_.accept(checker)
+    parse_result.checker_accept_all(checker)
 
     start = TestArgument.from_fn_def(fixture_def, checker=checker, source="test")
     assert start is not None