-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathcodegen.py
More file actions
279 lines (222 loc) · 9.25 KB
/
codegen.py
File metadata and controls
279 lines (222 loc) · 9.25 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
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
import json
from pathlib import Path
from jinja2 import Template
import re
# 获取项目根目录（脚本所在目录的父目录）
project_root = Path(__file__).parent.parent
# 定义模板和配置目录路径
# template_path = project_root / "codegen" / "templates" / "control_api.jinja2"
configs_dir = project_root / "codegen" / "configs"
print("Project root:", project_root)
CODE_GEN_HINT = """This file is auto generated by the code generation script.
Do not modify this file manually.
Use the `make codegen` command to regenerate.
当前文件为自动生成的控制 API 客户端代码。请勿手动修改此文件。
使用 `make codegen` 命令重新生成。"""
def with_hint(content, source="") -> str:
hint = CODE_GEN_HINT
if source:
hint += f"\n\nsource: {source}"
if content.startswith('"""'):
return f'"""\n{hint}\n\n{content[3:]}'
else:
return f'"""\n{hint}\n"""\n\n{content}'
def generate_jinja2(specific_config=None):
"""生成 Jinja2 模板代码
Args:
specific_config: 可选的特定配置文件路径，如果提供则只处理该文件
"""
def generate(config, source):
# 渲染模板
template_path = config["template"]
with open(
configs_dir / ".." / "templates" / template_path, "r", encoding="utf-8"
) as f:
template = Template(f.read())
rendered_code = template.render(**config)
# 确定输出路径（相对于项目根）
output_path = project_root / config["output_path"]
# 确保输出目录存在
output_path.parent.mkdir(parents=True, exist_ok=True)
# 写入生成的代码
import os
if os.path.exists(output_path):
os.chmod(output_path, 0o644)
with open(output_path, "w", encoding="utf-8") as f:
f.write(with_hint(rendered_code, source))
print(f"Generated {output_path}")
if specific_config:
# 只处理指定的配置文件
config_file = Path(specific_config)
if not config_file.is_absolute():
config_file = project_root / config_file
if config_file.suffix == ".yaml":
import yaml
with open(config_file, "r", encoding="utf-8") as f:
config = yaml.safe_load(f)
else:
with open(config_file, "r", encoding="utf-8") as f:
config = json.load(f)
generate(config, config_file.relative_to(project_root))
else:
# 处理所有配置文件
for config_file in configs_dir.glob("*.json"):
# 加载 JSON 配置
with open(config_file, "r", encoding="utf-8") as f:
config = json.load(f)
generate(config, config_file.relative_to(project_root))
for config_file in configs_dir.glob("*.yaml"):
# 加载 YAML 配置
import yaml
with open(config_file, "r", encoding="utf-8") as f:
config = yaml.safe_load(f)
generate(config, config_file.relative_to(project_root))
def generate_sync_code(specific_template=None):
"""生成同步代码
Args:
specific_template: 可选的特定模板文件名，如果提供则只处理匹配的文件
"""
if specific_template:
# 只处理特定的模板文件 - 直接构造完整路径
template_path = project_root / specific_template
if template_path.exists():
_generate_sync_code_for_file(template_path)
else:
print(f"Template file not found: {template_path}")
else:
# 处理所有模板文件
for async_file in project_root.glob("**/__*_async_template.py"):
_generate_sync_code_for_file(async_file)
def _generate_sync_code_for_file(async_file):
async_code = ""
with open(async_file, "r", encoding="utf-8") as f:
async_code = f.read()
# 找到所有的函数段落
results = []
in_async_def_indent = -1
async_block = []
pending_decorators = [] # 待处理的装饰器（还不知道是async函数还是sync函数）
for line in async_code.split("\n"):
strip_line = line.strip()
# 计算当前行的缩进级别
level = 0
for char in line:
if char == " ":
level += 1
else:
break
# 检测装饰器（只在不在 async 函数内部时收集）
if strip_line.startswith("@"):
# 如果当前在 async 函数内，且遇到了缩进级别 <= 函数定义级别的装饰器
# 说明退出了当前 async 函数
if in_async_def_indent >= 0 and level <= in_async_def_indent:
# 退出 async 函数
in_async_def_indent = -1
results.extend(async_block)
async_block = []
if in_async_def_indent == -1:
# 不在 async 函数内，收集装饰器
pending_decorators.append(line)
continue # 不立即添加到 results，等函数定义时再处理
# 检测函数定义
if (
strip_line.startswith("async def ")
or strip_line.startswith("def ")
or strip_line.startswith("class ")
):
if level <= in_async_def_indent:
# 退出 async 函数
in_async_def_indent = -1
results.extend(async_block)
async_block = []
if in_async_def_indent == -1:
# 未在 async 函数内
if strip_line.startswith("async def "):
# 这是一个 async 函数
in_async_def_indent = level
# 1. 先把装饰器添加到 results（原始的 async 版本）
results.extend(pending_decorators)
# 2. 把装饰器的 sync 版本添加到 async_block
for decorator_line in pending_decorators:
async_block.append(
decorator_line.replace("_async(", "(")
.replace("_async", "_sync")
.replace("Async", "Sync")
)
else:
# 这是一个 sync 函数，直接把装饰器添加到 results
results.extend(pending_decorators)
# 清空待处理的装饰器
pending_decorators = []
# 如果有待处理的装饰器，但遇到了非装饰器、非函数定义的行
# 说明这些装饰器后面不是函数，应该添加到 results
if (
pending_decorators
and strip_line
and not strip_line.startswith("@")
and not strip_line.startswith("async def ")
and not strip_line.startswith("def ")
):
results.extend(pending_decorators)
pending_decorators = []
# 如果在 async 函数内，添加 sync 版本到 async_block
if in_async_def_indent >= 0:
content = (
line.replace("AsyncClient", "Client")
.replace("AsyncOTSClient", "OTSClient")
.replace("AsyncOpenAI", "OpenAI")
.replace("AsyncMemory", "Memory")
.replace("async_playwright", "sync_playwright")
.replace("asyncio.gather(*", "(")
.replace("asynchronously", "synchronously")
.replace("_async", "")
.replace("async def", "def")
.replace("await ", "")
.replace("Async", "Sync")
.replace("__aenter__", "__enter__")
.replace("__aexit__", "__exit__")
.replace("asyncio.sleep", "time.sleep")
.replace("import asyncio", "")
.replace("async with", "with")
.replace("异步", "同步")
)
content = re.sub(
r"Awaitable\[(.*)\]",
"\\1",
content,
)
async_block.append(content)
# 添加原始行到 results
results.append(line)
results.extend(async_block)
sync_code = "\n".join(results)
sync_path = (
project_root
/ async_file.parent
/ async_file.name.replace("_async_template", "").replace("__", "")
)
import os
if os.path.exists(sync_path):
os.chmod(sync_path, 0o644)
with open(sync_path, "w", encoding="utf-8") as f:
f.write(with_hint(sync_code, str(async_file.relative_to(project_root))))
print(f"Generated {sync_path}")
def main():
import argparse
parser = argparse.ArgumentParser(description="代码生成器")
parser.add_argument("--sync-only", action="store_true", help="只生成同步代码")
parser.add_argument(
"--jinja2-only", action="store_true", help="只生成 Jinja2 模板代码"
)
parser.add_argument("--template", help="只处理指定的模板文件")
parser.add_argument("--config", help="只处理指定的配置文件")
args = parser.parse_args()
if args.jinja2_only:
generate_jinja2(args.config)
elif args.sync_only:
generate_sync_code(args.template)
else:
generate_jinja2()
generate_sync_code(args.template)
if __name__ == "__main__":
main()