use-walrus-if codemod can handle unused variables (#477)

* use-walrus-if codemod can handle unused variables

* cleanup

Author: clavedeluna

src/core_codemods/use_walrus_if.py

@@ -6,6 +6,7 @@
 from libcst._position import CodeRange
 from libcst.metadata import ParentNodeProvider, ScopeProvider
 
+from codemodder.codemods.utils_mixin import NameResolutionMixin
 from core_codemods.api import Metadata, Reference, ReviewGuidance, SimpleCodemod
 
 FoundAssign = namedtuple("FoundAssign", ["assign", "target", "value"])
@@ -17,7 +18,7 @@ def pairwise(iterable):
     return zip(a, b)
 
 
-class UseWalrusIf(SimpleCodemod):
+class UseWalrusIf(SimpleCodemod, NameResolutionMixin):
     metadata = Metadata(
         name="use-walrus-if",
         summary="Use Assignment Expression (Walrus) In Conditional",
@@ -76,6 +77,16 @@ def _filter_if(self, node: cst.CSTNode) -> cst.BaseExpression | None:
                 return test
         return None
 
+    def _single_access(self, original_node: cst.IfExp) -> bool:
+        match original_node.test:
+            case cst.Name():
+                access = self.find_accesses(original_node.test)
+            case cst.UnaryOperation():
+                access = self.find_accesses(original_node.test.expression)
+            case _:
+                access = self.find_accesses(original_node.test.left)
+        return len(access) == 1
+
     def on_visit(self, node: cst.CSTNode) -> Optional[bool]:
         if len(node.children) < 2:
             return super().on_visit(node)
@@ -87,7 +98,6 @@ def on_visit(self, node: cst.CSTNode) -> Optional[bool]:
                 continue
 
             assign, target, value = found_assign
-
             match if_test:
                 # If test can be a comparison expression
                 case cst.Comparison(
@@ -130,16 +140,24 @@ def leave_If(self, original_node, updated_node):
         if (result := self._if_stack.pop()) is not None:
             position, named_expr = result
             self.add_change_from_position(position, self.change_description)
+
+            # If a variable has a single access, it means it's only assigned and not used again.
+            # In this case, do not use a walrus named expr to prevent unused variable warnings.
+            # Instead, move the variable's rhs directly into the if statement.
+            new_expression = (
+                named_expr.value if self._single_access(original_node) else named_expr
+            )
+
             match updated_node.test:
                 case cst.Name():
-                    return updated_node.with_changes(test=named_expr)
+                    return updated_node.with_changes(test=new_expression)
                 case cst.UnaryOperation():
                     return updated_node.with_changes(
-                        test=updated_node.test.with_changes(expression=named_expr)
+                        test=updated_node.test.with_changes(expression=new_expression)
                     )
                 case _:
                     return updated_node.with_changes(
-                        test=updated_node.test.with_changes(left=named_expr)
+                        test=updated_node.test.with_changes(left=new_expression)
                     )
 
         return original_node

tests/codemods/test_walrus_if.py

@@ -225,3 +225,47 @@ def test_multiple_walrus_changes(self, tmpdir):
         assert changes[0].lineNumber == 2
         assert changes[1].lineNumber == 5
         assert changes[2].lineNumber == 8
+
+    def test_no_walrus_if_unused_variable(self, tmpdir):
+        input_code = """
+        test = (
+            1
+            if True
+            else 0
+        )
+        if test:
+            print("var will not be reused")
+        """
+        expected_output = """
+        if (
+            1
+            if True
+            else 0
+        ):
+            print("var will not be reused")
+        """
+        self.run_and_assert(tmpdir, input_code, expected_output)
+
+    def test_no_walrus_if_not(self, tmpdir):
+        input_code = """
+        val = do_something()
+        if not val:
+            print("something")
+        """
+        expected_output = """
+        if not do_something():
+            print("something")
+        """
+        self.run_and_assert(tmpdir, input_code, expected_output)
+
+    def test_no_walrus_if(self, tmpdir):
+        input_code = """
+        val = do_something()
+        if val is None:
+            print("hi")
+        """
+        expected_output = """
+        if do_something() is None:
+            print("hi")
+        """
+        self.run_and_assert(tmpdir, input_code, expected_output)