Skip to content

Instantly share code, notes, and snippets.

@sangwoo-joh
Last active June 14, 2022 03:56
Show Gist options
  • Save sangwoo-joh/26e9007ebc2de256b0b3deeda051772d to your computer and use it in GitHub Desktop.
Save sangwoo-joh/26e9007ebc2de256b0b3deeda051772d to your computer and use it in GitHub Desktop.
A rough example code for renaming variables using libcst.
import libcst
import difflib
from termcolor import colored
class RenameTransformer(libcst.CSTTransformer):
def __init__(self, rename_pairs):
self.rename_pairs = rename_pairs
self.restore_keywords = []
def _rename(self, original_node, renamed_node):
if original_node.value in self.rename_pairs.keys():
return renamed_node.with_changes(value=self.rename_pairs[original_node.value])
else:
return renamed_node
def leave_Name(self, original_node, renamed_node):
"""
LibCST parses variables in f-strings.
All [Name] nodes are already parsed correctly.
"""
return self._rename(original_node, renamed_node)
def visit_Arg(self, node):
"""
LibCST does not provide APIs to decide whether to visit the child attributes or not.
Some tricky hack is needed, e.g. restoring the original keyword as follow.
"""
if node.keyword and node.keyword.value in self.rename_pairs.keys():
self.restore_keywords.append(node.keyword.value)
return True
def leave_Arg(self, original_node, renamed_node):
try:
restore = self.restore_keywords.pop()
return renamed_node.with_changes(keyword=renamed_node.keyword.with_changes(value=restore))
except IndexError:
# stack is empty
return renamed_node
test_code = """
def foo(bar, zar):
print(f'bar is {bar}, not {zar}')
print('bar is {bar}, not {zar}'.format(bar=bar, zar=zar))
"""
if __name__ == '__main__':
rename_candidates = {
'bar': 'bar_new',
}
rename_transformer = RenameTransformer(rename_candidates)
original_tree = libcst.parse_module(test_code)
print(original_tree)
renamed_tree = original_tree.visit(rename_transformer)
print('#' * 100)
print('original code:')
print(test_code)
print('#' * 100)
print('renaming bar -> bar_new:')
# restoring rename tree as code is easy
diff = "".join(difflib.unified_diff(test_code.splitlines(1), renamed_tree.code.splitlines(1)))
for l in diff.splitlines():
if l.startswith('+'):
print(colored(l, 'green'))
elif l.startswith('-'):
print(colored(l, 'red'))
else:
print(l)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment