Last active
June 14, 2022 03:56
-
-
Save sangwoo-joh/26e9007ebc2de256b0b3deeda051772d to your computer and use it in GitHub Desktop.
A rough example code for renaming variables using libcst.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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