diff --git a/src/codemodder/codemods/use_walrus_if.py b/src/codemodder/codemods/use_walrus_if.py index 27dba7032..930006679 100644 --- a/src/codemodder/codemods/use_walrus_if.py +++ b/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 diff --git a/tests/codemods/test_walrus_if.py b/tests/codemods/test_walrus_if.py index c1b8c9e9f..7c9b7a7ce 100644 --- a/tests/codemods/test_walrus_if.py +++ b/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) """