Add --unsafe

Refs #84

Author: akx

auto_walrus.py

@@ -34,6 +34,7 @@
 @dataclasses.dataclass
 class Config:
     line_length: int
+    unsafe: bool = False
 
 
 def name_lineno_coloffset_iterable(
@@ -199,6 +200,7 @@ def related_vars_are_unused(
 
 def visit_function_def(
     node: ast.FunctionDef,
+    config: Config,
 ) -> list[tuple[Token, Token]]:
     names = set()
     assignments: set[Token] = set()
@@ -210,7 +212,7 @@ def visit_function_def(
     related_vars: dict[str, list[Token]] = {}
     in_body_vars: dict[Token, set[Token]] = {}
 
-    for _node in node.body:
+    for _node in ast.walk(node) if config.unsafe else node.body:
         if isinstance(_node, ast.Assign):
             process_assign(_node, assignments, related_vars)
         elif isinstance(_node, ast.If):
@@ -276,7 +278,7 @@ def auto_walrus(
     walruses = []
     for node in ast.walk(tree):
         if isinstance(node, ast.FunctionDef):
-            walruses.extend(visit_function_def(node))
+            walruses.extend(visit_function_def(node, config))
     lines_to_remove = []
     walruses = sorted(walruses, key=lambda x: (-x[1][1], -x[1][2]))
 
@@ -371,6 +373,11 @@ def main(argv: Sequence[str] | None = None) -> int:  # pragma: no cover
         required=False,
         default=r"^$",
     )
+    parser.add_argument(
+        "--unsafe",
+        action="store_true",
+        help="Also process if statements inside other blocks (like for loops)",
+    )
     # black formatter's default
     parser.add_argument("--line-length", type=int, default=88)
     args = parser.parse_args(argv)
@@ -402,7 +409,10 @@ def main(argv: Sequence[str] | None = None) -> int:  # pragma: no cover
                     content = fd.read()
             except UnicodeDecodeError:
                 continue
-            new_content = auto_walrus(content, args.line_length)
+            new_content = auto_walrus(
+                content,
+                Config(line_length=args.line_length, unsafe=args.unsafe),
+            )
             if new_content is not None and content != new_content:
                 sys.stdout.write(f"Rewriting {filepath}\n")
                 with open(filepath, "w", encoding="utf-8") as fd:

tests/main_test.py

@@ -7,6 +7,7 @@
 
 import pytest
 
+from auto_walrus import Config
 from auto_walrus import auto_walrus
 from auto_walrus import main
 
@@ -129,6 +130,20 @@ def test_noop(src: str) -> None:
     assert ret is None
 
 
+@pytest.mark.parametrize(
+    ("src", "expected"),
+    [
+        (
+            'def foo(data):\n    if True:\n        foo = data.get("blah")\n        if foo:\n            return foo\n    return data',
+            'def foo(data):\n    if True:\n        if (foo := data.get("blah")):\n            return foo\n    return data',
+        ),
+    ],
+)
+def test_rewrite_unsafe(src: str, expected: str) -> None:
+    ret = auto_walrus(src, Config(line_length=88, unsafe=True))
+    assert ret == expected
+
+
 ProjectDirT = Tuple[pathlib.Path, List[pathlib.Path]]
 
 SRC_ORIG = "def foo():\n    a = 0\n    if a:\n        print(a)\n"