-
Notifications
You must be signed in to change notification settings - Fork 2
/
model_transformer.py
70 lines (49 loc) · 1.94 KB
/
model_transformer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
import libcst as cst
from typing import List
class ModelTransformer(cst.CSTTransformer):
def __init__(self, updated_classes: List[cst.ClassDef]):
self.updated_classes = updated_classes
def on_visit(self, node: cst.CSTNode) -> bool:
return True
def on_leave(self, original_node: cst.CSTNode, updated_node: cst.CSTNode) -> cst.CSTNode:
#print(f"Visiting node: {original_node}")
#print(f"Length of updated_nodes: {len(self.updated_nodes)}")
if len(self.updated_classes) == 0:
return updated_node
if isinstance(original_node, cst.ClassDef):
print("Found a class definition: ", original_node.name)
removal_index = -1
for i, node in enumerate(self.updated_classes):
#print(type(node))
#print(node)
if node.name.value == original_node.name.value:
#print(f"Found a match: {node.name.value} == {original_node.name.value}")
removal_index = i
updated_node = node
break
if removal_index != -1:
self.updated_classes.pop(removal_index)
else:
pass
return updated_node
def add_module(to_add: List[cst.ClassDef], module_to_update: cst.Module):
# Append the new modules to the end of the existing module
new_body = list(module_to_update.body)
for new_module in to_add:
new_body.append(new_module)
module_to_update = module_to_update.with_changes(body=tuple(new_body))
return module_to_update
def apply_codemod_to_file(modules_to_add: List[cst.Module]):
source_code = ""
with open("generated_models.py", "r") as file:
source_code = file.read()
# Parse the source code into a CST
tree = cst.parse_module(source_code)
# Apply the codemod
transformer = ModelTransformer(updated_nodes=modules_to_add)
modified_tree = tree.visit(transformer)
modified_tree = add_module(modules_to_add, modified_tree)
#print(f"\n\n\n\nModified code:\n{modified_tree.code}")
with open("generated_models.py", "w") as file:
file.write(modified_tree.code)
file.close()