Add support for methods with func type comment excluding self/cls (#622)

* Add support for methods with func type comment excluding self/cls

PEP 484 doesn't really specify carefully how function type
comments should work on methods, but since usually the type of
`self` / `cls` is automatic, most use cases choose to only annotate
the other arguments.

As a result, this commit modifies our codemod so that non-static
methods can specify either all the arguments, or all but one of
them. We'll correctly zip together the inline and func-type-comment
types either way, typically getting no type for `cls` or `self`.

We accomplish this by using matchers to trigger the visit
method for FunctionDef rather than using visit_FunctionDef, which gives
us enough context to determine when a function def is a regular
function versus a method (plus also matching the decorators against
`@staticmethod`, so that we trigger the normal function logic in
that case).

Co-authored-by: Zsolt Dollenstein <[email protected]>

libcst/codemod/commands/convert_type_comments.py

@@ -13,6 +13,7 @@
 from typing_extensions import TypeAlias
 
 import libcst as cst
+import libcst.matchers as m
 from libcst.codemod import CodemodContext, VisitorBasedCodemodCommand
 
 
@@ -285,6 +286,7 @@ def is_empty(self) -> bool:
     def from_cst(
         cls,
         node_cst: cst.FunctionDef,
+        is_method: bool,
     ) -> "FunctionTypeInfo":
         """
         Using the `ast` type comment extraction logic, get type information
@@ -350,6 +352,18 @@ def from_cst(
                     },
                     returns=returns,
                 )
+            elif is_method and len(argtypes) == len(args) - 1:
+                # Merge as above, but skip merging the initial `self` or `cls` arg.
+                return cls(
+                    arguments={
+                        args[0].arg: args[0].type_comment,
+                        **{
+                            arg.arg: arg.type_comment or ast.unparse(from_func_type)
+                            for arg, from_func_type in zip(args[1:], argtypes)
+                        },
+                    },
+                    returns=returns,
+                )
             else:
                 # On arity mismatches, ignore the type information
                 return cls({}, None)
@@ -623,10 +637,16 @@ def leave_With(
     # (B) we also manually reach down to the first statement inside of the
     #     funciton body and aggressively strip type comments from leading
     #     whitespaces
+    #
+    # PEP 484 underspecifies how to apply type comments to (non-static)
+    # methods - it would be possible to provide a type for `self`, or to omit
+    # it. So we accept either approach when interpreting type comments on
+    # non-static methods: the first argument an have a type provided or not.
 
-    def visit_FunctionDef(
+    def _visit_FunctionDef(
         self,
         node: cst.FunctionDef,
+        is_method: bool,
     ) -> None:
         """
         Set up the data we need to handle function definitions:
@@ -636,11 +656,35 @@ def visit_FunctionDef(
         - Set that we are aggressively stripping type comments, which will
           remain true until we visit the body.
         """
-        function_type_info = FunctionTypeInfo.from_cst(node)
+        function_type_info = FunctionTypeInfo.from_cst(node, is_method=is_method)
         self.aggressively_strip_type_comments = not function_type_info.is_empty()
         self.function_type_info_stack.append(function_type_info)
         self.function_body_stack.append(node.body)
 
+    @m.call_if_not_inside(m.ClassDef())
+    @m.visit(m.FunctionDef())
+    def visit_method(
+        self,
+        node: cst.FunctionDef,
+    ) -> None:
+        return self._visit_FunctionDef(
+            node=node,
+            is_method=False,
+        )
+
+    @m.call_if_inside(m.ClassDef())
+    @m.visit(m.FunctionDef())
+    def visit_function(
+        self,
+        node: cst.FunctionDef,
+    ) -> None:
+        return self._visit_FunctionDef(
+            node=node,
+            is_method=not any(
+                m.matches(d.decorator, m.Name("staticmethod")) for d in node.decorators
+            ),
+        )
+
     def leave_TrailingWhitespace(
         self,
         original_node: cst.TrailingWhitespace,

libcst/codemod/commands/tests/test_convert_type_comments.py

@@ -337,6 +337,73 @@ def f(
         """
         self.assertCodemod39Plus(before, after)
 
+    def test_method_transforms(self) -> None:
+        before = """
+        class A:
+
+            def __init__(self, thing):  # type: (str) -> None
+                self.thing = thing
+
+            @classmethod
+            def make(cls):  # type: () -> A
+                return cls("thing")
+
+            @staticmethod
+            def f(x, y):  # type: (object, object) -> None
+                pass
+
+            def method0(
+                self,
+                other_thing,
+            ):  # type: (str) -> bool
+                return self.thing == other_thing
+
+            def method1(
+                self,  # type: A
+                other_thing,  # type: str
+            ):  # type: (int) -> bool
+                return self.thing == other_thing
+
+            def method2(
+                self,
+                other_thing,
+            ):  # type: (A, str) -> bool
+                return self.thing == other_thing
+        """
+        after = """
+        class A:
+
+            def __init__(self, thing: str) -> None:
+                self.thing = thing
+
+            @classmethod
+            def make(cls) -> "A":
+                return cls("thing")
+
+            @staticmethod
+            def f(x: object, y: object) -> None:
+                pass
+
+            def method0(
+                self,
+                other_thing: str,
+            ) -> bool:
+                return self.thing == other_thing
+
+            def method1(
+                self: "A",
+                other_thing: str,
+            ) -> bool:
+                return self.thing == other_thing
+
+            def method2(
+                self: "A",
+                other_thing: str,
+            ) -> bool:
+                return self.thing == other_thing
+        """
+        self.assertCodemod39Plus(before, after)
+
     def test_no_change_if_function_type_comments_unused(self) -> None:
         before = """
         # arity error in arguments