Skip to content

Commit

Permalink
Preserve comments from removed lines
Browse files Browse the repository at this point in the history
  • Loading branch information
drdavella committed Sep 21, 2023
1 parent 888b437 commit 0ab5033
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 2 deletions.
15 changes: 13 additions & 2 deletions src/codemodder/codemods/use_walrus_if.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,11 +109,22 @@ def leave_Assign(self, original_node, updated_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.
Preserves the whitespace and comments in the line 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)
trailing_whitespace = (
(
original_node.trailing_whitespace.with_changes(
whitespace=cst.SimpleWhitespace(""),
),
)
if original_node.trailing_whitespace.comment
else ()
)
return cst.FlattenSentinel(
original_node.leading_lines + trailing_whitespace
)

return updated_node
13 changes: 13 additions & 0 deletions tests/codemods/test_walrus_if.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,19 @@ def test_walrus_if_name_only(self, tmpdir):
"""
self.run_and_assert(tmpdir, input_code, expected_output)

def test_walrus_if_preserve_comments(self, tmpdir):
input_code = """
val = do_something() # comment
if val is not None: # another comment
do_something_else(val)
"""
expected_output = """
# comment
if (val := do_something()) is not None: # another comment
do_something_else(val)
"""
self.run_and_assert(tmpdir, input_code, expected_output)

def test_walrus_if_multiple(self, tmpdir):
input_code = """
val = do_something()
Expand Down

0 comments on commit 0ab5033

Please sign in to comment.