-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathinclude_task.py
More file actions
135 lines (107 loc) · 4.1 KB
/
include_task.py
File metadata and controls
135 lines (107 loc) · 4.1 KB
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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
import os
import json
import re
def browse_py_files(folder):
"""
Browse all the .py files in the specified folder.
"""
py_files = []
for file in os.listdir(folder):
if file.endswith('.py'):
py_files.append(file)
if '__init__.py' in py_files:
py_files.remove('__init__.py')
return py_files
def get_urdfs_and_description(py_file):
"""
Extract URDFs and task description from the Python file.
"""
urdfs = []
description = ''
with open(py_file, 'r') as f:
lines = f.readlines()
in_method = False
for line in lines:
# Extract task description
if 'class' in line and description == '':
# Look for a comment right under the class definition
if lines[lines.index(line) + 1].strip().startswith('"""'):
for next_line in lines[lines.index(line) + 1:]:
if lines.index(next_line) > lines.index(line) + 1:
description += ' '
description += next_line.strip()
if next_line.strip().endswith('"""'):
description = description[3:-3]
break
else:
# If no comment found, refer to self.lang_template in __init__() method
for next_line in lines[lines.index(line) + 1:]:
if 'self.lang_template' in next_line:
description = next_line.split('=')[-1].strip().strip('"')
break
# Extract URDFs
# Search for method definition
if 'def reset(' in line:
in_method = True
# Extract URDFs within method body
if in_method:
urdf_matches = re.findall(r"'(.*?\.urdf)'", line)
urdfs.extend(urdf_matches)
return urdfs, description
def update_generated_tasks(py_files, json_file):
"""
Update the generated_tasks.json file with information from Python files.
"""
with open(json_file, 'r') as f:
data = json.load(f)
# Convert the list to a set for faster lookup
registered_files = set(data.keys())
for py_file in py_files:
# Convert '_' to '-' and remove '.py' extension for key
key = py_file.replace('_', '-').replace('.py', '')
if key not in registered_files:
urdfs, description = get_urdfs_and_description(
os.path.join(r'./cliport/generated_tasks', py_file))
# Format URDFs list
# formatted_urdfs = []
# for urdf in urdfs:
# formatted_urdfs.append(f'{urdf}')
# Update JSON data
data[key] = {
"task-name": key,
"task-description": description,
"assets-used": urdfs
}
# Write updated data back to JSON file
with open(json_file, 'w') as f:
json.dump(data, f, indent=4)
def update_json(file_list, json_file):
"""
Update the JSON file with the list of Python files.
"""
with open(json_file, 'r') as f:
data = json.load(f)
# Convert the list to a set for faster lookup
registered_files = set(data)
# Check for missing files
missing_files = [file for file in file_list if file not in registered_files]
# Update JSON data
if missing_files:
data.extend(missing_files)
with open(json_file, 'w') as f:
json.dump(data, f, indent=4)
def main():
# Folder containing Python files
folder = r'./cliport/generated_tasks'
# JSON file containing registered file names
json_file = r'./prompts/data/generated_task_codes.json'
generated_tasks_json = r'./prompts/data/generated_tasks.json'
# Browse .py files
py_files = browse_py_files(folder)
# Update JSON file
update_json(py_files, json_file)
# Update generated_tasks.json file
update_generated_tasks(py_files, generated_tasks_json)
print("Updated JSON file with missing Python files.")
if __name__ == "__main__":
main()