-
Notifications
You must be signed in to change notification settings - Fork 0
/
Working with AST.py
100 lines (78 loc) · 3.33 KB
/
Working with AST.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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import os
import shutil
import ast
READ_MODE = "r"
WRITE_MODE = "w"
def process_file(path: str, filename: str) -> tuple:
"""
Process the given Python source file in the specified path.
Args:
path (str): The path where the source file is located.
filename (str): The name of the source file to process.
Returns:
tuple: A tuple indicating whether the process was successful (bool) and a message (str).
"""
try:
# Read the content of the source file
with open(os.path.join(path, filename), READ_MODE) as file:
source_code = file.read()
# Parse the source code into an AST
ast_tree = ast.parse(source_code)
# Traverse the AST and modify import statements
for node in ast.walk(ast_tree):
if isinstance(node, ast.ImportFrom):
if node.level > 1:
# Extract the folder containing the shared module
shared_module_folder = (
node.module.split(".")[-2]
if len(node.module.split(".")) > 1
else ""
)
# Create the folder if it doesn't exist
shared_module_path = os.path.join(path, shared_module_folder)
if not os.path.exists(shared_module_path):
os.mkdir(shared_module_path)
# Copy the shared module file to the shared module folder
module_filename = node.module.split(".")[-1] + ".py"
destination_file_path = os.path.join(
shared_module_path, module_filename
)
shutil.copyfile(
node.module.replace(".", os.sep) + ".py", destination_file_path
)
# Update the import level to 1
node.level = 1
# Generate modified source code from the modified AST
modified_source_code = ast.unparse(ast_tree)
# Write the modified source code back to the file
with open(os.path.join(path, filename), WRITE_MODE) as file:
file.write(modified_source_code)
except Exception as e:
# Return error message if an exception occurs
return (True, str(e))
# Return success message if the process completes without errors
return (False, "Successfully completed!")
def main():
"""
Main function to iterate over files in the directory and process each Python file.
"""
# Iterate over files in the directory and process each Python file
is_error = False
for root, dirs, files in os.walk("./"):
if is_error:
break
for filename in files:
# Skip files not under the 'proj' directory
if not root.startswith("./proj"):
continue
# Process the file and print the result
is_error, result = process_file(root, filename)
if is_error:
print(f"ERROR: {result}")
break
else:
print(f"RESULT: {os.path.join(root, filename)} - {result}")
# Check if the script is executed as the main module
if __name__ == "__main__":
main()
# %%