125 lines
4.3 KiB
Python
125 lines
4.3 KiB
Python
# generate_pyi.py
|
||
|
||
import re
|
||
import os
|
||
import argparse
|
||
from pathlib import Path
|
||
|
||
def sanitize_name(name: str) -> str:
|
||
name = re.sub(r'[^a-zA-Z0-9_]', '_', name)
|
||
return re.sub(r'^[^a-zA-Z_]+', '', name)
|
||
|
||
def pascal_case(s: str) -> str:
|
||
"""将字符串转为 PascalCase"""
|
||
return ''.join(word.capitalize() for word in re.split(r'[\W_]+', s) if word)
|
||
|
||
def generate_pyi(base_dir: str, output_file: str):
|
||
base_path = Path(base_dir)
|
||
if not base_path.exists():
|
||
raise FileNotFoundError(f"资源目录不存在: {base_dir}")
|
||
|
||
# 构建目录树结构
|
||
tree = {}
|
||
|
||
for root, dirs, files in os.walk(str(base_path)):
|
||
relative_path = Path(root).relative_to(base_path)
|
||
parts = list(relative_path.parts) if relative_path != Path('.') else []
|
||
|
||
# 插入当前路径到树中
|
||
current = tree
|
||
for part in parts:
|
||
safe_part = sanitize_name(part)
|
||
if safe_part not in current:
|
||
current[safe_part] = {}
|
||
current = current[safe_part]
|
||
|
||
# 添加图像文件作为叶子节点
|
||
for file in files:
|
||
if file.startswith('.'):
|
||
continue
|
||
if file.split('.')[-1].lower() in {'png', 'jpg', 'jpeg', 'bmp', 'gif'}:
|
||
stem = sanitize_name(Path(file).stem)
|
||
current[stem] = None # 表示这是一个图像资源（str 类型）
|
||
|
||
# 构建类定义
|
||
class_defs = []
|
||
written_classes = set()
|
||
|
||
def build_class(class_name: str, data: dict):
|
||
if not data or class_name in written_classes:
|
||
return
|
||
written_classes.add(class_name)
|
||
|
||
lines = [f"class {class_name}:\n"]
|
||
has_content = False
|
||
|
||
for key in sorted(data.keys()):
|
||
value = data[key]
|
||
safe_key = key
|
||
if value is None:
|
||
# 图像文件：直接定义为 str 属性
|
||
lines.append(f" {safe_key}: str\n")
|
||
has_content = True
|
||
else:
|
||
# 子目录：引用一个类
|
||
nested_class_name = pascal_case(safe_key)
|
||
lines.append(f" {safe_key}: \"{nested_class_name}\"\n")
|
||
build_class(nested_class_name, value)
|
||
has_content = True
|
||
|
||
if has_content:
|
||
class_defs.append("".join(lines))
|
||
else:
|
||
# 空类不写入
|
||
written_classes.remove(class_name)
|
||
|
||
# 构建主类 ImageResourceLoader
|
||
loader_lines = []
|
||
for key in sorted(tree.keys()):
|
||
value = tree[key]
|
||
if not value: # 跳过空的子目录
|
||
loader_lines.append(f" {key}: str\n")
|
||
continue
|
||
class_name = pascal_case(key)
|
||
loader_lines.append(f" {key}: {class_name}\n")
|
||
build_class(class_name, value)
|
||
|
||
# 写入 .pyi 文件
|
||
with open(output_file, 'w', encoding='utf-8') as f:
|
||
f.write("from typing import Optional, Any\n\n")
|
||
|
||
# 写入所有子类定义
|
||
for line in class_defs:
|
||
f.write(line + "\n")
|
||
|
||
# 写入主类
|
||
f.write("class ImageResourceLoader:\n")
|
||
f.write(" def __init__(self, base_dir: str): ...\n")
|
||
f.write(" def list_resources(self) -> None: ...\n")
|
||
f.write(" def get(self, path: str, default: Optional[str] = ...) -> Optional[str]: ...\n")
|
||
f.write(" def __getattr__(self, name: str) -> Any: ...\n")
|
||
f.write("\n # 自动生成的资源类属性（仅用于 IDE 提示）\n")
|
||
f.writelines(loader_lines)
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(description="为 ImageResourceLoader 类生成 .pyi 存根文件。")
|
||
parser.add_argument("source_file", help="Python 源文件路径（例如：image_loader.py）")
|
||
parser.add_argument("--base_dir", default=".automation",
|
||
help="图像资源目录，默认值：.automation")
|
||
args = parser.parse_args()
|
||
|
||
source_path = Path(args.source_file)
|
||
if not source_path.exists():
|
||
raise FileNotFoundError(f"源文件不存在: {args.source_file}")
|
||
|
||
# 输出 .pyi 文件路径
|
||
pyi_path = source_path.with_suffix(".pyi")
|
||
|
||
print(f"生成存根文件: {pyi_path}")
|
||
print(f"使用资源目录: {args.base_dir}")
|
||
|
||
generate_pyi(args.base_dir, str(pyi_path))
|
||
"""
|
||
python scripts/generate_pyi.py image_loader.py --base_dir assets/images
|
||
|
||
""" |