Skip -> None on abstract and stub methods (#34)

JelleZijlstra · web-flow · commit d30ed2eb2d1f · 2022-08-21T14:19:36.000-07:00
diff --git a/README.md b/README.md
@@ -80,6 +80,8 @@ Things to add:
 Unreleased
 
 - Add `--pyanalyze-report`
+- Do not add `None` return types to methods marked with `@abstractmethod` and
+  to methods in stub files
 
 21.12.0 (December 21, 2021)
 
diff --git a/autotyping/autotyping.py b/autotyping/autotyping.py
@@ -1,5 +1,6 @@
 import argparse
 from dataclasses import dataclass, field
+import enum
 import json
 from typing import Dict, List, Optional, Sequence, Set, Tuple, Type
 from typing_extensions import TypedDict
@@ -35,6 +36,11 @@ class PyanalyzeSuggestion(TypedDict):
     imports: List[str]
 
 
+class DecoratorKind(enum.Enum):
+    asynq = 1
+    abstractmethod = 2
+
+
 @dataclass
 class State:
     annotate_optionals: List[NamedParam]
@@ -85,6 +91,8 @@ class AutotypeCommand(VisitorBasedCodemodCommand):
     DESCRIPTION: str = "Automatically adds simple type annotations."
     METADATA_DEPENDENCIES = (PositionProvider,)
 
+    state: State
+
     @staticmethod
     def add_args(arg_parser: argparse.ArgumentParser) -> None:
         arg_parser.add_argument(
@@ -246,6 +254,10 @@ def __init__(
             only_without_imports=only_without_imports,
         )
 
+    def is_stub(self) -> bool:
+        filename = self.context.filename
+        return filename is not None and filename.endswith(".pyi")
+
     def visit_FunctionDef(self, node: libcst.FunctionDef) -> None:
         self.state.seen_return_statement.append(False)
         self.state.seen_raise_statement.append(False)
@@ -277,7 +289,9 @@ def leave_Lambda(
     def leave_FunctionDef(
         self, original_node: libcst.FunctionDef, updated_node: libcst.FunctionDef
     ) -> libcst.CSTNode:
-        is_asynq = any(is_asynq_decorator(dec) for dec in original_node.decorators)
+        kinds = {get_decorator_kind(decorator) for decorator in updated_node.decorators}
+        is_asynq = DecoratorKind.asynq in kinds
+        is_abstractmethod = DecoratorKind.abstractmethod in kinds
         seen_return = self.state.seen_return_statement.pop()
         seen_raise = self.state.seen_raise_statement.pop()
         seen_yield = self.state.seen_yield.pop()
@@ -338,6 +352,8 @@ def leave_FunctionDef(
             and not seen_raise
             and not seen_return
             and (is_asynq or not seen_yield)
+            and not is_abstractmethod
+            and not self.is_stub()
         ):
             return updated_node.with_changes(
                 returns=libcst.Annotation(annotation=libcst.Name(value="None"))
@@ -555,16 +571,26 @@ def type_of_expression(expr: libcst.BaseExpression) -> Optional[Type[object]]:
         return None
 
 
-def is_asynq_decorator(dec: libcst.Decorator) -> bool:
-    """Is this @asynq()?"""
-    if not isinstance(dec.decorator, libcst.Call):
-        return False
-    call = dec.decorator
-    if not isinstance(call.func, libcst.Name):
-        return False
-    if call.func.value != "asynq":
-        return False
-    if call.args:
-        # @asynq() with custom arguments may do something unexpected
-        return False
-    return True
+def get_decorator_kind(dec: libcst.Decorator) -> Optional[DecoratorKind]:
+    """Is this @asynq() or @abstractmethod?"""
+    if isinstance(dec.decorator, libcst.Call):
+        call = dec.decorator
+        if not isinstance(call.func, libcst.Name):
+            return None
+        if call.func.value != "asynq":
+            return None
+        if call.args:
+            # @asynq() with custom arguments may do something unexpected
+            return None
+        return DecoratorKind.asynq
+    elif isinstance(dec.decorator, libcst.Name):
+        if dec.decorator.value == "abstractmethod":
+            return DecoratorKind.abstractmethod
+    elif isinstance(dec.decorator, libcst.Attribute):
+        if (
+            dec.decorator.attr.value == "abstractmethod"
+            and isinstance(dec.decorator.value, libcst.Name)
+            and dec.decorator.value.value == "abc"
+        ):
+            return DecoratorKind.abstractmethod
+    return None
diff --git a/tests/test_codemod.py b/tests/test_codemod.py
@@ -1,4 +1,4 @@
-from libcst.codemod import CodemodTest
+from libcst.codemod import CodemodTest, CodemodContext
 from autotyping.autotyping import AutotypeCommand
 
 
@@ -28,6 +28,14 @@ def bar():
 
             def baz():
                 return
+
+            @abstractmethod
+            def very_abstract():
+                pass
+
+            @abc.abstractmethod
+            def very_abstract_without_import_from():
+                pass
         """
         after = """
             def foo() -> None:
@@ -38,9 +46,33 @@ def bar():
 
             def baz() -> None:
                 return
+
+            @abstractmethod
+            def very_abstract():
+                pass
+
+            @abc.abstractmethod
+            def very_abstract_without_import_from():
+                pass
         """
         self.assertCodemod(before, after, none_return=True)
 
+    def test_none_return_stub(self) -> None:
+        before = """
+            def foo():
+                pass
+        """
+        after = """
+            def foo():
+                pass
+        """
+        self.assertCodemod(
+            before,
+            after,
+            none_return=True,
+            context_override=CodemodContext(filename="stub.pyi"),
+        )
+
     def test_scalar_return(self) -> None:
         before = """
             def foo():
