Add --safe and --aggressive (#37)

JelleZijlstra · web-flow · commit 1b905a0bdb2e · 2022-09-05T15:38:42.000-07:00
Fixes #36
diff --git a/README.md b/README.md
@@ -70,15 +70,20 @@ more transformations. The following are supported:
     new imports. This is useful because suggestions that require imports may need
     more manual work.
 
-Things to add:
+There are two shortcut flags to enable multiple transformations at once:
 
-- Infer `-> bool` as the return type if all return statements are
-  boolean expressions like `==`.
+- `--safe` enables changes that should always be safe. This includes
+  `--none-return`, `--scalar-return`, and `--annotate-magics`.
+- `--aggressive` enables riskier changes that are more likely to produce
+  new type checker errors. It includes all of `--safe` as well as `--bool-param`,
+  `--int-param`, `--float-param`, `--str-param`, `--bytes-param`, and
+  `--annotate-imprecise-magics`.
 
 # Changelog
 
 Unreleased
 
+- Add `--safe` and `--aggressive`
 - Add `--pyanalyze-report`
 - Do not add `None` return types to methods marked with `@abstractmethod` and
   to methods in stub files
diff --git a/autotyping/autotyping.py b/autotyping/autotyping.py
@@ -85,6 +85,36 @@ class State:
 }
 
 
+class _SafeAction(argparse.Action):
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,
+        namespace: argparse.Namespace,
+        values: object,
+        option_string: Optional[str] = ...,
+    ) -> None:
+        namespace.none_return = True
+        namespace.scalar_return = True
+        namespace.annotate_magics = True
+
+
+class _AggressiveAction(_SafeAction):
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,
+        namespace: argparse.Namespace,
+        values: object,
+        option_string: Optional[str] = ...,
+    ) -> None:
+        super().__call__(parser, namespace, values, option_string)
+        namespace.bool_param = True
+        namespace.int_param = True
+        namespace.float_param = True
+        namespace.str_param = True
+        namespace.bytes_param = True
+        namespace.annotate_imprecise_magics = True
+
+
 class AutotypeCommand(VisitorBasedCodemodCommand):
 
     # Add a description so that future codemodders can see what this does.
@@ -187,6 +217,18 @@ def add_args(arg_parser: argparse.ArgumentParser) -> None:
             default=False,
             help="Only apply pyanalyze suggestions that do not require imports",
         )
+        arg_parser.add_argument(
+            "--safe",
+            action=_SafeAction,
+            help="Apply all safe transformations",
+            nargs="?",
+        )
+        arg_parser.add_argument(
+            "--aggressive",
+            action=_AggressiveAction,
+            help="Apply all transformations that do not require arguments",
+            nargs="?",
+        )
 
     def __init__(
         self,
@@ -205,6 +247,8 @@ def __init__(
         int_param: bool = False,
         pyanalyze_report: Optional[str] = None,
         only_without_imports: bool = False,
+        safe: bool = False,
+        aggressive: bool = False,
     ) -> None:
         super().__init__(context)
         param_type_pairs = [
