Work around issue in libcst to preserve whitespace

src/codemodder/codemods/use_walrus_if.py

@@ -106,3 +106,14 @@ def leave_Assign(self, original_node, updated_node):
                 return cst.RemoveFromParent()
 
         return original_node
+
+    def leave_SimpleStatementLine(self, original_node, updated_node):
+        """
+        This is a workaround for the fact that libcst doesn't preserve the whitespace in the parent node when all children are removed.
+
+        This feels like a bug in libCST but we'll work around it for now.
+        """
+        if not updated_node.body:
+            return cst.FlattenSentinel(original_node.leading_lines)
+
+        return updated_node

tests/codemods/test_walrus_if.py

@@ -50,10 +50,10 @@ def test_walrus_if_multiple(self, tmpdir):
 if foo == "bar":
     whatever(foo)
 """
-        # TODO: not sure why libcst isn't preserving empty lines
         expected_output = """
 if (val := do_something()) is not None:
     do_something_else(val)
+
 if (foo := hello()) == "bar":
     whatever(foo)
 """