-
Notifications
You must be signed in to change notification settings - Fork 394
Expand file tree
/
Copy pathcommon_utils.py
More file actions
235 lines (184 loc) · 6.71 KB
/
Copy pathcommon_utils.py
File metadata and controls
235 lines (184 loc) · 6.71 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
"""通用工具函数模块,提供任务 ID 生成、文件操作和文档转换等功能。"""
import os
import shutil
import datetime
import hashlib
import tomllib
from app.schemas.enums import CompTemplate
from app.utils.log_util import logger
import re
import pypandoc # type: ignore[import-unresolved]
from app.config.setting import settings
TASK_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$")
def create_task_id() -> str:
"""生成基于时间戳和随机哈希的唯一任务 ID。"""
# 生成时间戳和随机hash
timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
random_hash = hashlib.md5(str(datetime.datetime.now()).encode()).hexdigest()[:8]
return f"{timestamp}-{random_hash}"
def ensure_safe_task_id(task_id: str) -> str:
"""验证任务 ID 的合法性,防止路径遍历攻击。
Args:
task_id: 待验证的任务 ID。
Returns:
验证通过的任务 ID。
Raises:
ValueError: 任务 ID 不合法时抛出。
"""
normalized = (task_id or "").strip()
if not normalized or not TASK_ID_PATTERN.fullmatch(normalized):
raise ValueError("非法 task_id")
return normalized
def create_work_dir(task_id: str) -> str:
"""为指定任务创建工作目录,并复制字体文件到工作目录。
Args:
task_id: 任务 ID。
Returns:
工作目录路径。
"""
# 设置主工作目录和子目录
work_dir = os.path.join("project", "work_dir", task_id)
try:
# 创建目录,如果目录已存在也不会报错
os.makedirs(work_dir, exist_ok=True)
# 复制字体文件到工作目录,确保图表中文正常显示
_copy_fonts_to_work_dir(work_dir)
return work_dir
except Exception as e:
# 捕获并记录创建目录时的异常
logger.error(f"创建工作目录失败: {str(e)}")
raise
# 字体源目录(backend/fonts/)
_FONTS_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "fonts")
def _copy_fonts_to_work_dir(work_dir: str) -> None:
"""将后端字体目录中的字体文件复制到工作目录。
Args:
work_dir: 目标工作目录路径。
"""
fonts_dir = os.path.normpath(_FONTS_DIR)
if not os.path.isdir(fonts_dir):
logger.warning(f"字体目录不存在: {fonts_dir}")
return
for filename in os.listdir(fonts_dir):
if not filename.lower().endswith((".ttf", ".otf", ".ttc")):
continue
src = os.path.join(fonts_dir, filename)
dst = os.path.join(work_dir, filename)
try:
shutil.copy2(src, dst)
logger.debug(f"复制字体: {filename} -> {work_dir}")
except Exception as e:
logger.warning(f"复制字体 {filename} 失败: {e}")
def get_work_dir(task_id: str) -> str:
"""获取指定任务的工作目录路径。
Args:
task_id: 任务 ID。
Returns:
工作目录路径。
Raises:
FileNotFoundError: 工作目录不存在时抛出。
"""
work_dir = os.path.join("project", "work_dir", task_id)
if os.path.exists(work_dir):
return work_dir
else:
logger.error(f"工作目录不存在: {work_dir}")
raise FileNotFoundError(f"工作目录不存在: {work_dir}")
# TODO: 是不是应该将 Prompt 写成一个 class
def get_config_template(comp_template: CompTemplate = CompTemplate.CHINA) -> dict:
"""获取论文模板配置。
Args:
comp_template: 竞赛模板类型。
Returns:
模板配置字典。
"""
if comp_template == CompTemplate.CHINA:
return load_toml(os.path.join("app", "config", "md_template.toml"))
return {}
def load_toml(path: str) -> dict:
"""加载 TOML 配置文件。
Args:
path: TOML 文件路径。
"""
with open(path, "rb") as f:
return tomllib.load(f)
def load_markdown(path: str) -> str:
"""加载 Markdown 文件内容。
Args:
path: Markdown 文件路径。
"""
with open(path, "r", encoding="utf-8") as f:
return f.read()
def get_current_files(folder_path: str, type: str = "all") -> list[str]:
"""获取指定目录下的文件列表。
Args:
folder_path: 目录路径。
type: 文件类型过滤(all/md/ipynb/data/image)。
"""
files = os.listdir(folder_path)
if type == "all":
return files
elif type == "md":
return [file for file in files if file.endswith(".md")]
elif type == "ipynb":
return [file for file in files if file.endswith(".ipynb")]
elif type == "data":
return [
file for file in files if file.endswith(".xlsx") or file.endswith(".csv")
]
elif type == "image":
return [
file for file in files if file.endswith(".png") or file.endswith(".jpg")
]
return []
def transform_link(task_id: str, content: str):
"""将 Markdown 中的图片链接转换为静态资源 URL。
Args:
task_id: 任务 ID,用于构建 URL 路径。
content: 包含图片链接的 Markdown 文本。
"""
content = re.sub(
r"!\[(.*?)\]\((.*?\.(?:png|jpg|jpeg|gif|bmp|webp))\)",
lambda match: f"})",
content,
)
return content
def md_2_docx(task_id: str):
"""将 Markdown 论文转换为 DOCX 格式。
Args:
task_id: 任务 ID。
"""
work_dir = get_work_dir(task_id)
md_path = os.path.join(work_dir, "res.md")
docx_path = os.path.join(work_dir, "res.docx")
extra_args = [
"--resource-path",
str(work_dir),
"--mathml", # MathML 格式公式
"--standalone",
]
pypandoc.convert_file(
source_file=md_path,
to="docx",
outputfile=docx_path,
format="markdown+tex_math_dollars",
extra_args=extra_args,
)
print(f"转换完成: {docx_path}")
logger.info(f"转换完成: {docx_path}")
def split_footnotes(text: str) -> tuple[str, list[tuple[str, str]]]:
"""从文本中分离正文和脚注。
Args:
text: 包含脚注的完整文本。
Returns:
(正文, 脚注列表) 的元组,脚注格式为 (编号, 内容)。
"""
main_text = re.sub(
r"\n\[\^\d+\]:.*?(?=\n\[\^|\n\n|\Z)", "", text, flags=re.DOTALL
)
# 正文中的脚注引用标记(如 [^1])本身不是可读内容,随定义一起去除
main_text = re.sub(r"\[\^\d+\]", "", main_text).strip()
# 匹配脚注定义
footnotes = re.findall(r"\[\^(\d+)\]:\s*(.+?)(?=\n\[\^|\n\n|\Z)", text, re.DOTALL)
logger.info(f"main_text:{main_text} \n footnotes:{footnotes}")
return main_text, footnotes