Merge pull request #90 from akx/avoid-duplicate-walrus-removal

Fix incorrect rewrite with --unsafe on nested functions

Author: MarcoGorelli

auto_walrus.py

@@ -271,12 +271,12 @@ def auto_walrus(
     except SyntaxError:  # pragma: no cover
         return None
 
-    walruses = []
+    walrus_set: set[tuple[Token, Token]] = set()
     for node in ast.walk(tree):
         if isinstance(node, ast.FunctionDef):
-            walruses.extend(visit_function_def(node, config))
+            walrus_set.update(visit_function_def(node, config))
     lines_to_remove = []
-    walruses = sorted(walruses, key=lambda x: (-x[1][1], -x[1][2]))
+    walruses = sorted(walrus_set, key=lambda x: (-x[1][1], -x[1][2]))
 
     if not walruses:
         return None

tests/main_test.py

@@ -137,6 +137,20 @@ def test_noop(src: str) -> None:
             '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',
         ),
+        # Nested functions - issue #89
+        (
+            "def foo():\n"
+            "    def bar():\n"
+            "        def quox():\n"
+            "            conn_time_zone = fetch_rel_time_zone(df.native)\n"
+            "            if conn_time_zone != time_zone:\n"
+            "                print(conn_time_zone)\n",
+            "def foo():\n"
+            "    def bar():\n"
+            "        def quox():\n"
+            "            if (conn_time_zone := fetch_rel_time_zone(df.native)) != time_zone:\n"
+            "                print(conn_time_zone)\n",
+        ),
     ],
 )
 def test_rewrite_unsafe(src: str, expected: str) -> None: