Add tests for __exit__ inference

Author: chadrik

mypy/stubgen.py

@@ -119,6 +119,7 @@
     fail_missing,
     find_module_path_and_all_py3,
     generate_guarded,
+    infer_method_arg_types,
     infer_method_ret_type,
     remove_misplaced_type_comments,
     report_missing,
@@ -480,7 +481,7 @@ def _get_func_args(self, o: FuncDef, ctx: FunctionContext) -> list[ArgSig]:
             # name their 0th argument other than self/cls
             is_self_arg = i == 0 and name == "self"
             is_cls_arg = i == 0 and name == "cls"
-            typename = ""
+            typename: str | None = None
             if annotated_type and not is_self_arg and not is_cls_arg:
                 # Luckily, an argument explicitly annotated with "Any" has
                 # type "UnboundType" and will not match.
@@ -500,6 +501,15 @@ def _get_func_args(self, o: FuncDef, ctx: FunctionContext) -> list[ArgSig]:
 
             args.append(ArgSig(name, typename, default=bool(arg_.initializer)))
 
+        if ctx.class_info is not None and all(
+            arg.type is None and arg.default is False for arg in args
+        ):
+            new_args = infer_method_arg_types(
+                ctx.name, ctx.class_info.self_var, [arg.name for arg in args]
+            )
+            if new_args is not None:
+                args = new_args
+
         is_dataclass_generated = (
             self.analyzed and self.processing_dataclass and o.info.names[o.name].plugin_generated
         )

mypy/stubgenc.py

@@ -33,6 +33,7 @@
     ClassInfo,
     FunctionContext,
     SignatureGenerator,
+    infer_method_arg_types,
     infer_method_ret_type,
 )
 
@@ -251,7 +252,7 @@ def get_default_function_sig(self, func: object, ctx: FunctionContext) -> Functi
                 # method:
                 return FunctionSig(
                     name=ctx.name,
-                    args=infer_method_args(ctx.name, ctx.class_info.self_var),
+                    args=infer_c_method_args(ctx.name, ctx.class_info.self_var),
                     ret_type=infer_method_ret_type(ctx.name),
                 )
             else:
@@ -306,6 +307,16 @@ def get_annotation(key: str) -> str | None:
         if kwargs:
             arglist.append(ArgSig(f"**{kwargs}", get_annotation(kwargs)))
 
+        # add types for known special methods
+        if ctx.class_info is not None and all(
+            arg.type is None and arg.default is False for arg in arglist
+        ):
+            new_args = infer_method_arg_types(
+                ctx.name, ctx.class_info.self_var, [arg.name for arg in arglist if arg.name]
+            )
+            if new_args is not None:
+                arglist = new_args
+
         ret_type = get_annotation("return") or infer_method_ret_type(ctx.name)
         return FunctionSig(ctx.name, arglist, ret_type)
 
@@ -829,7 +840,9 @@ def is_pybind_skipped_attribute(attr: str) -> bool:
     return attr.startswith("__pybind11_module_local_")
 
 
-def infer_method_args(name: str, self_var: str = "self") -> list[ArgSig]:
+def infer_c_method_args(
+    name: str, self_var: str = "self", arg_names: list[str] | None = None
+) -> list[ArgSig]:
     args: list[ArgSig] | None = None
     if name.startswith("__") and name.endswith("__"):
         name = name[2:-2]
@@ -928,6 +941,10 @@ def infer_method_args(name: str, self_var: str = "self") -> list[ArgSig]:
                 ArgSig(name="value", type="BaseException | None"),
                 ArgSig(name="traceback", type="types.TracebackType | None"),
             ]
+    if args is None:
+        args = infer_method_arg_types(name, self_var, arg_names)
+    else:
+        args = [ArgSig(name=self_var)] + args
     if args is None:
         args = [ArgSig(name="*args"), ArgSig(name="**kwargs")]
-    return [ArgSig(name=self_var or "self")] + args
+    return args

mypy/stubutil.py

@@ -16,7 +16,7 @@
 import mypy.options
 from mypy.modulefinder import ModuleNotFoundReason
 from mypy.moduleinspect import InspectError, ModuleInspect
-from mypy.stubdoc import FunctionSig
+from mypy.stubdoc import ArgSig, FunctionSig
 from mypy.types import AnyType, NoneType, Type, TypeList, TypeStrVisitor, UnboundType, UnionType
 
 # Modules that may fail when imported, or that may have side effects (fully qualified).
@@ -312,6 +312,7 @@ def fullname(self) -> str:
 
 
 def infer_method_ret_type(name: str) -> str | None:
+    """Infer return types for known special methods"""
     if name.startswith("__") and name.endswith("__"):
         name = name[2:-2]
         if name in ("float", "bool", "bytes", "int", "complex", "str"):
@@ -328,6 +329,34 @@ def infer_method_ret_type(name: str) -> str | None:
     return None
 
 
+def infer_method_arg_types(
+    name: str, self_var: str = "self", arg_names: list[str] | None = None
+) -> list[ArgSig] | None:
+    """Infer argument types for known special methods"""
+    args: list[ArgSig] | None = None
+    if name.startswith("__") and name.endswith("__"):
+        if arg_names and len(arg_names) >= 1 and arg_names[0] == "self":
+            arg_names = arg_names[1:]
+
+        name = name[2:-2]
+        if name == "exit":
+            if arg_names is None:
+                arg_names = ["type", "value", "traceback"]
+            if len(arg_names) == 3:
+                arg_types = [
+                    "type[BaseException] | None",
+                    "BaseException | None",
+                    "types.TracebackType | None",
+                ]
+                args = [
+                    ArgSig(name=arg_name, type=arg_type)
+                    for arg_name, arg_type in zip(arg_names, arg_types)
+                ]
+    if args is not None:
+        return [ArgSig(name=self_var)] + args
+    return None
+
+
 @mypyc_attr(allow_interpreted_subclasses=True)
 class SignatureGenerator:
     """Abstract base class for extracting a list of FunctionSigs for each function."""

mypy/test/teststubgen.py

@@ -33,7 +33,7 @@
     mypy_options,
     parse_options,
 )
-from mypy.stubgenc import InspectionStubGenerator, infer_method_args
+from mypy.stubgenc import InspectionStubGenerator, infer_c_method_args
 from mypy.stubutil import (
     ClassInfo,
     common_dir_prefix,
@@ -785,36 +785,36 @@ class StubgencSuite(unittest.TestCase):
     """
 
     def test_infer_hash_sig(self) -> None:
-        assert_equal(infer_method_args("__hash__"), [self_arg])
+        assert_equal(infer_c_method_args("__hash__"), [self_arg])
         assert_equal(infer_method_ret_type("__hash__"), "int")
 
     def test_infer_getitem_sig(self) -> None:
-        assert_equal(infer_method_args("__getitem__"), [self_arg, ArgSig(name="index")])
+        assert_equal(infer_c_method_args("__getitem__"), [self_arg, ArgSig(name="index")])
 
     def test_infer_setitem_sig(self) -> None:
         assert_equal(
-            infer_method_args("__setitem__"),
+            infer_c_method_args("__setitem__"),
             [self_arg, ArgSig(name="index"), ArgSig(name="object")],
         )
         assert_equal(infer_method_ret_type("__setitem__"), "None")
 
     def test_infer_eq_op_sig(self) -> None:
         for op in ("eq", "ne", "lt", "le", "gt", "ge"):
             assert_equal(
-                infer_method_args(f"__{op}__"), [self_arg, ArgSig(name="other", type="object")]
+                infer_c_method_args(f"__{op}__"), [self_arg, ArgSig(name="other", type="object")]
             )
 
     def test_infer_binary_op_sig(self) -> None:
         for op in ("add", "radd", "sub", "rsub", "mul", "rmul"):
-            assert_equal(infer_method_args(f"__{op}__"), [self_arg, ArgSig(name="other")])
+            assert_equal(infer_c_method_args(f"__{op}__"), [self_arg, ArgSig(name="other")])
 
     def test_infer_equality_op_sig(self) -> None:
         for op in ("eq", "ne", "lt", "le", "gt", "ge", "contains"):
             assert_equal(infer_method_ret_type(f"__{op}__"), "bool")
 
     def test_infer_unary_op_sig(self) -> None:
         for op in ("neg", "pos"):
-            assert_equal(infer_method_args(f"__{op}__"), [self_arg])
+            assert_equal(infer_c_method_args(f"__{op}__"), [self_arg])
 
     def test_infer_cast_sig(self) -> None:
         for op in ("float", "bool", "bytes", "int"):

test-data/unit/stubgen.test

@@ -3600,6 +3600,38 @@ class Some:
     def __float__(self) -> float: ...
     def __index__(self) -> int: ...
 
+
+[case testKnownMagicMethodsArgTypes]
+class MismatchNames:
+    def __exit__(self, tp, val, tb): ...
+
+class MatchNames:
+    def __exit__(self, type, value, traceback): ...
+
+[out]
+class MismatchNames:
+    def __exit__(self, tp: type[BaseException] | None, val: BaseException | None, tb: types.TracebackType | None) -> None: ...
+
+class MatchNames:
+    def __exit__(self, type: type[BaseException] | None, value: BaseException | None, traceback: types.TracebackType | None) -> None: ...
+
+-- Same as above (but can generate import statements)
+[case testKnownMagicMethodsArgTypes_inspect]
+class MismatchNames:
+    def __exit__(self, tp, val, tb): ...
+
+class MatchNames:
+    def __exit__(self, type, value, traceback): ...
+
+[out]
+import types
+
+class MismatchNames:
+    def __exit__(self, tp: type[BaseException] | None, val: BaseException | None, tb: types.TracebackType | None): ...
+
+class MatchNames:
+    def __exit__(self, type: type[BaseException] | None, value: BaseException | None, traceback: types.TracebackType | None): ...
+
 [case testTypeVarPEP604Bound]
 from typing import TypeVar
 T = TypeVar("T", bound=str | None)