Fix the --safe argument (#52)

JelleZijlstra · web-flow · commit ad9433c82685 · 2023-01-20T05:00:24.000-08:00
Fixes #44
diff --git a/README.md b/README.md
@@ -82,19 +82,28 @@ There are two shortcut flags to enable multiple transformations at once:
 
 # Changelog
 
-22.9.0 (September 5, 2022)
+## Unreleased
+
+- Add `--guess-common-names` (contributed by John Litborn)
+- Fix the `--safe` and `--aggressive` flags so they don't take
+  ignored arguments
+- `--length-hint` should return `int` (contributed by Nikita Sobolev)
+- Fix bug in import adding (contributed by Shantanu)
+
+## 22.9.0 (September 5, 2022)
 
 - Add `--safe` and `--aggressive`
 - Add `--pyanalyze-report`
 - Do not add `None` return types to methods marked with `@abstractmethod` and
   to methods in stub files
 - Improve type inference:
+
   - `"string" % ...` is always `str`
   - `b"bytes" % ...` is always `bytes`
   - An `and` or `or` operator where left and right sides are of the same type
     returns that type
   - `is`, `is not`, `in`, and `not in` always return `bool`
 
-21.12.0 (December 21, 2021)
+## 21.12.0 (December 21, 2021)
 
 - Initial PyPI release
diff --git a/autotyping/autotyping.py b/autotyping/autotyping.py
@@ -4,7 +4,6 @@
 import json
 from typing import Dict, List, Optional, Sequence, Set, Tuple, Type
 from typing_extensions import TypedDict
-import re
 
 import libcst
 from libcst.codemod import CodemodContext, VisitorBasedCodemodCommand
@@ -90,37 +89,6 @@ 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
-        namespace.guess_common_names = True
-
-
 class AutotypeCommand(VisitorBasedCodemodCommand):
     # Add a description so that future codemodders can see what this does.
     DESCRIPTION: str = "Automatically adds simple type annotations."
@@ -230,15 +198,15 @@ def add_args(arg_parser: argparse.ArgumentParser) -> None:
         )
         arg_parser.add_argument(
             "--safe",
-            action=_SafeAction,
+            action="store_true",
+            default=False,
             help="Apply all safe transformations",
-            nargs="?",
         )
         arg_parser.add_argument(
             "--aggressive",
-            action=_AggressiveAction,
+            action="store_true",
+            default=False,
             help="Apply all transformations that do not require arguments",
-            nargs="?",
         )
 
     def __init__(
@@ -262,6 +230,18 @@ def __init__(
         safe: bool = False,
         aggressive: bool = False,
     ) -> None:
+        if safe or aggressive:
+            none_return = True
+            scalar_return = True
+            annotate_magics = True
+        if aggressive:
+            bool_param = True
+            int_param = True
+            float_param = True
+            str_param = True
+            bytes_param = True
+            annotate_imprecise_magics = True
+            guess_common_names = True
         super().__init__(context)
         param_type_pairs = [
             (bool_param, bool),
@@ -295,12 +275,16 @@ def __init__(
                     )
                 ] = metadata
         self.state = State(
-            annotate_optionals=[NamedParam.make(s) for s in annotate_optional]
-            if annotate_optional
-            else [],
-            annotate_named_params=[NamedParam.make(s) for s in annotate_named_param]
-            if annotate_named_param
-            else [],
+            annotate_optionals=(
+                [NamedParam.make(s) for s in annotate_optional]
+                if annotate_optional
+                else []
+            ),
+            annotate_named_params=(
+                [NamedParam.make(s) for s in annotate_named_param]
+                if annotate_named_param
+                else []
+            ),
             none_return=none_return,
             scalar_return=scalar_return,
             param_types={typ for param, typ in param_type_pairs if param},
