docparser_core/plugins/
subtitle.rs

1//! 字幕解析：srt / vtt / ass / ssa（零第三方依赖，语义对齐 Python `plugins/subtitle.py`）。
2//!
3//! 字幕不是纯文本：序号、`-->` 时间码、`{\pos(..)}` 覆盖块、VTT 内联标签都会污染下游
4//! （TTS 会念出时间码），所以单独一个后端，而不是落到 txt。
5//!
6//! 约定（两载体逐字一致，见 FORMAT_TODO P3）：
7//! * **一条字幕条目 = 一页**：`Page.name` = 规范化时间轴 `HH:MM:SS,mmm --> HH:MM:SS,mmm`
8//!   （VTT 的 `.`、ASS 的厘秒都归一到这里），文本是该页唯一的 Paragraph 块，
9//!   条目内换行保留为 `\n`。
10//! * 文本清理：先剥 `<...>` 标签，再解 6 个实体（`&amp; &lt; &gt; &quot; &apos; &nbsp;`），
11//!   两者顺序不能反，否则 `&lt;b&gt;` 会被当成标签吞掉。
12//! * ASS/SSA 额外：删 `{...}` 覆盖块，`\N` / `\n` → 换行，`\h` → 空格。
13//! * 清理后为空的条目跳过；一条不剩则 `corrupt`。
14//! * 严格 UTF-8（可带 BOM）：猜编码会把中文念成乱码，宁可报错让调用方先转码。
15
16use std::sync::Arc;
17
18use crate::error::{DocError, Result};
19use crate::model::{Block, BlockKind, Document, Page};
20use crate::plugin::{FormatPlugin, ParseCapability, RawDoc, RawPage, RawTextItem};
21use crate::registry::Registry;
22
23/// 支持的扩展名（`ssa` 与 `ass` 共用同一套事件表解析）。
24pub const FORMATS: &[&str] = &["srt", "vtt", "ass", "ssa"];
25
26/// 实体名 → 字符；未知实体（`&#123;` 等）原样保留。
27fn entity(name: &str) -> Option<char> {
28    match name {
29        "amp" => Some('&'),
30        "lt" => Some('<'),
31        "gt" => Some('>'),
32        "quot" => Some('"'),
33        "apos" => Some('\''),
34        "nbsp" => Some(' '),
35        _ => None,
36    }
37}
38
39/// 剥离 `<...>` 标签；`<` 后没有 `>` 时按普通字符保留（不吞掉半个文本）。
40pub fn strip_tags(text: &str) -> String {
41    let chars: Vec<char> = text.chars().collect();
42    let mut out = String::with_capacity(text.len());
43    let mut i = 0;
44    while i < chars.len() {
45        if chars[i] == '<' {
46            if let Some(off) = chars[i + 1..].iter().position(|c| *c == '>') {
47                i += off + 2;
48                continue;
49            }
50        }
51        out.push(chars[i]);
52        i += 1;
53    }
54    out
55}
56
57/// 只解 6 个常见实体；`&` 后 10 个字符内找不到 `;` 就按普通字符处理。
58pub fn decode_entities(text: &str) -> String {
59    let chars: Vec<char> = text.chars().collect();
60    let mut out = String::with_capacity(text.len());
61    let mut i = 0;
62    while i < chars.len() {
63        if chars[i] == '&' {
64            if let Some(off) = chars[i + 1..].iter().position(|c| *c == ';') {
65                if off <= 9 {
66                    let name: String = chars[i + 1..i + 1 + off].iter().collect();
67                    if let Some(c) = entity(&name) {
68                        out.push(c);
69                        i += off + 2;
70                        continue;
71                    }
72                }
73            }
74        }
75        out.push(chars[i]);
76        i += 1;
77    }
78    out
79}
80
81/// 标签 → 实体 → 逐行去空白；保留条目内换行。
82pub fn clean_text(raw: &str) -> String {
83    let body = decode_entities(&strip_tags(raw));
84    body.split('\n')
85        .map(|line| line.trim())
86        .collect::<Vec<_>>()
87        .join("\n")
88        .trim()
89        .to_string()
90}
91
92/// 删除 ASS 覆盖块 `{...}`（`{\pos(..)}` / `{\i1}` 等）。
93pub fn strip_ass_overrides(text: &str) -> String {
94    let chars: Vec<char> = text.chars().collect();
95    let mut out = String::with_capacity(text.len());
96    let mut i = 0;
97    while i < chars.len() {
98        if chars[i] == '{' {
99            if let Some(off) = chars[i + 1..].iter().position(|c| *c == '}') {
100                i += off + 2;
101                continue;
102            }
103        }
104        out.push(chars[i]);
105        i += 1;
106    }
107    out
108}
109
110/// ASS 文本：去覆盖块 → `\N` / `\n` 转换行 → `\h` 转空格。
111pub fn ass_text(raw: &str) -> String {
112    strip_ass_overrides(raw)
113        .replace("\\N", "\n")
114        .replace("\\n", "\n")
115        .replace("\\h", " ")
116}
117
118fn all_digits(s: &str) -> bool {
119    !s.is_empty() && s.bytes().all(|b| b.is_ascii_digit())
120}
121
122/// `[H:]MM:SS[.,]fff` → 毫秒；无法解析返回 `None`。
123///
124/// 小数位不足 3 位右补零（`.5` → 500ms，ASS 厘秒 `.50` → 500ms），多出 3 位截断。
125pub fn parse_clock(spec: &str) -> Option<i64> {
126    let s = spec.trim();
127    if s.is_empty() {
128        return None;
129    }
130    let parts: Vec<&str> = s.split(':').collect();
131    if parts.len() != 2 && parts.len() != 3 {
132        return None;
133    }
134    let (head, tail) = parts.split_at(parts.len() - 1);
135    if !head.iter().all(|p| all_digits(p)) {
136        return None;
137    }
138    let tail = tail[0];
139    let sep = if tail.contains(',') {
140        ','
141    } else if tail.contains('.') {
142        '.'
143    } else {
144        return None;
145    };
146    let (sec, frac) = tail.split_once(sep)?;
147    if !all_digits(sec) || !all_digits(frac) {
148        return None;
149    }
150    let mut milli = frac.chars().take(3).collect::<String>();
151    while milli.len() < 3 {
152        milli.push('0');
153    }
154    let milli: i64 = milli.parse().ok()?;
155    let (hours, minutes) = if head.len() == 2 {
156        (head[0].parse::<i64>().ok()?, head[1].parse::<i64>().ok()?)
157    } else {
158        (0, head[0].parse::<i64>().ok()?)
159    };
160    let total = hours
161        .saturating_mul(3600)
162        .saturating_add(minutes.saturating_mul(60))
163        .saturating_add(sec.parse::<i64>().ok()?);
164    Some(total.saturating_mul(1000).saturating_add(milli))
165}
166
167/// 毫秒 → `HH:MM:SS,mmm`（小时不截断，不足两位补零）。
168pub fn format_clock(ms: i64) -> String {
169    let ms = if ms < 0 { 0 } else { ms };
170    let milli = ms % 1000;
171    let total = ms / 1000;
172    let second = total % 60;
173    let minutes_total = total / 60;
174    let minute = minutes_total % 60;
175    let hour = minutes_total / 60;
176    format!("{hour:02}:{minute:02}:{second:02},{milli:03}")
177}
178
179/// `start --> end [settings]` → (start_ms, end_ms)；非法返回 `None`。
180pub fn parse_timeline(line: &str) -> Option<(i64, i64)> {
181    let pieces: Vec<&str> = line.split("-->").collect();
182    if pieces.len() != 2 {
183        return None;
184    }
185    let start = parse_clock(pieces[0])?;
186    let end_field = pieces[1].split_whitespace().next()?;
187    let end = parse_clock(end_field)?;
188    Some((start, end))
189}
190
191fn timeline_name(tl: (i64, i64)) -> String {
192    format!("{} --> {}", format_clock(tl.0), format_clock(tl.1))
193}
194
195/// 按空行分组；组内保留原始行（行尾空白去掉，组尾空行丢弃）。
196pub fn split_blocks(text: &str) -> Vec<Vec<String>> {
197    let mut blocks: Vec<Vec<String>> = Vec::new();
198    let mut buf: Vec<String> = Vec::new();
199    for line in text.split('\n') {
200        if !line.trim().is_empty() {
201            buf.push(line.trim_end().to_string());
202        } else if !buf.is_empty() {
203            blocks.push(std::mem::take(&mut buf));
204        }
205    }
206    if !buf.is_empty() {
207        blocks.push(buf);
208    }
209    blocks
210}
211
212fn first_timeline_index(block: &[String]) -> Option<usize> {
213    block.iter().position(|ln| ln.contains("-->"))
214}
215
216/// SRT：组内首个含 `-->` 的行是时间轴，其后全是文本；无时间轴的组跳过。
217pub fn parse_srt(text: &str) -> Result<Vec<(String, String)>> {
218    let mut cues = Vec::new();
219    for block in split_blocks(text) {
220        let Some(idx) = first_timeline_index(&block) else {
221            continue;
222        };
223        let Some(tl) = parse_timeline(&block[idx]) else {
224            let shown: String = block[idx].trim().chars().take(120).collect();
225            return Err(DocError::corrupt(format!("SRT 时间轴格式非法: {shown}")));
226        };
227        cues.push((timeline_name(tl), block[idx + 1..].join("\n")));
228    }
229    Ok(cues)
230}
231
232/// VTT：首行必须 `WEBVTT`；`NOTE` / `STYLE` / `REGION` 块跳过；时间轴后可带设置。
233pub fn parse_vtt(text: &str) -> Result<Vec<(String, String)>> {
234    let first = text
235        .split('\n')
236        .find(|ln| !ln.trim().is_empty())
237        .unwrap_or("");
238    if !first.trim().to_uppercase().starts_with("WEBVTT") {
239        return Err(DocError::corrupt("不是合法的 WebVTT（缺少 WEBVTT 头）"));
240    }
241    let mut cues = Vec::new();
242    for block in split_blocks(text) {
243        let head = block[0].trim();
244        let keyword = head.split(' ').next().unwrap_or("").to_uppercase();
245        if matches!(keyword.as_str(), "NOTE" | "STYLE" | "REGION") {
246            continue;
247        }
248        if head.to_uppercase().starts_with("WEBVTT") {
249            continue;
250        }
251        let Some(idx) = first_timeline_index(&block) else {
252            continue;
253        };
254        let Some(tl) = parse_timeline(&block[idx]) else {
255            let shown: String = block[idx].trim().chars().take(120).collect();
256            return Err(DocError::corrupt(format!("VTT 时间轴格式非法: {shown}")));
257        };
258        cues.push((timeline_name(tl), block[idx + 1..].join("\n")));
259    }
260    Ok(cues)
261}
262
263fn lower_ascii(s: &str) -> String {
264    s.to_ascii_lowercase()
265}
266
267/// ASS/SSA：只吃 `[Events]` 段的 `Format:` + `Dialogue:`（`Comment:` 跳过）。
268pub fn parse_ass(text: &str) -> Result<Vec<(String, String)>> {
269    let mut fields: Option<Vec<String>> = None;
270    let mut section = String::new();
271    let mut cues = Vec::new();
272
273    for line in text.split('\n') {
274        let row = line.trim();
275        if row.is_empty() {
276            continue;
277        }
278        if row.starts_with('[') && row.ends_with(']') {
279            section = lower_ascii(row[1..row.len() - 1].trim());
280            continue;
281        }
282        if section != "events" {
283            continue;
284        }
285        let low = lower_ascii(row);
286        if low.starts_with("format:") {
287            fields = Some(
288                row["format:".len()..]
289                    .split(',')
290                    .map(|f| lower_ascii(f.trim()))
291                    .collect(),
292            );
293            continue;
294        }
295        if low.starts_with("comment:") {
296            continue;
297        }
298        if !low.starts_with("dialogue:") {
299            continue;
300        }
301        let Some(names) = fields.as_ref() else {
302            return Err(DocError::corrupt("ASS 的 Dialogue 出现在 Format 行之前"));
303        };
304        let payload = row["dialogue:".len()..].trim_start();
305        // Text 是最后一个字段且可含逗号：限次拆分，剩余全部归最后一段
306        let values: Vec<&str> = payload.splitn(names.len(), ',').collect();
307        if values.len() != names.len() {
308            return Err(DocError::corrupt(format!(
309                "ASS Dialogue 字段数不足（期望 {}）",
310                names.len()
311            )));
312        }
313        let field = |name: &str| -> Option<&str> {
314            names
315                .iter()
316                .position(|f| f == name)
317                .map(|i| values[i].trim())
318        };
319        let (Some(start), Some(end)) = (
320            field("start").and_then(parse_clock),
321            field("end").and_then(parse_clock),
322        ) else {
323            let shown: String = payload.chars().take(120).collect();
324            return Err(DocError::corrupt(format!("ASS 时间轴格式非法: {shown}")));
325        };
326        cues.push((
327            timeline_name((start, end)),
328            ass_text(field("text").unwrap_or("")),
329        ));
330    }
331
332    if fields.is_none() {
333        return Err(DocError::corrupt("ASS 未找到 [Events] 段的 Format 行"));
334    }
335    Ok(cues)
336}
337
338/// 严格 UTF-8（可带 BOM）读取并归一换行；非法字节 → `corrupt`。
339pub fn read_subtitle(path: &str) -> Result<String> {
340    let raw =
341        std::fs::read(path).map_err(|e| DocError::corrupt(format!("无法读取字幕文件: {e}")))?;
342    let raw = raw.strip_prefix(&[0xEF, 0xBB, 0xBF]).unwrap_or(&raw);
343    let text = std::str::from_utf8(raw)
344        .map_err(|_| DocError::corrupt("字幕文件不是合法 UTF-8：请另存为 UTF-8 后重试"))?;
345    Ok(text.replace("\r\n", "\n").replace('\r', "\n"))
346}
347
348/// 解析成 `[(规范化时间轴, 已清理文本)]`，空白条目丢弃；一条不剩 → `corrupt`。
349pub fn parse_cues(path: &str, kind: &str) -> Result<Vec<(String, String)>> {
350    let text = read_subtitle(path)?;
351    let raw = match kind {
352        "vtt" => parse_vtt(&text)?,
353        "ass" | "ssa" => parse_ass(&text)?,
354        _ => parse_srt(&text)?,
355    };
356    let mut out: Vec<(String, String)> = Vec::new();
357    for (name, body) in raw {
358        let cleaned = clean_text(&body);
359        if !cleaned.is_empty() {
360            out.push((name, cleaned));
361        }
362    }
363    if out.is_empty() {
364        return Err(DocError::corrupt(format!("{kind} 未找到任何非空字幕条目")));
365    }
366    Ok(out)
367}
368
369/// 字幕解析后端（一个实现覆盖 srt/vtt/ass/ssa，靠 `kind` 分派）。
370pub struct SubtitleParseCapability {
371    pub id: &'static str,
372    pub kind: &'static str,
373}
374
375impl ParseCapability for SubtitleParseCapability {
376    fn id(&self) -> &str {
377        self.id
378    }
379    fn formats(&self) -> &[&str] {
380        match self.kind {
381            "vtt" => &["vtt"],
382            "ass" => &["ass", "ssa"],
383            _ => &["srt"],
384        }
385    }
386    fn parse(&self, path: &str) -> Result<RawDoc> {
387        let cues = parse_cues(path, self.kind)?;
388        let pages = cues
389            .into_iter()
390            .map(|(name, text)| RawPage {
391                name: Some(name),
392                items: vec![RawTextItem {
393                    text,
394                    ..Default::default()
395                }],
396                ..Default::default()
397            })
398            .collect();
399        Ok(RawDoc {
400            pages,
401            ..Default::default()
402        })
403    }
404}
405
406/// 字幕插件：每条一页，页内一个 Paragraph 块（与 Python `SubtitlePlugin` 同构）。
407pub struct SubtitlePlugin;
408
409impl FormatPlugin for SubtitlePlugin {
410    fn id(&self) -> &str {
411        "subtitle"
412    }
413    fn formats(&self) -> &[&str] {
414        FORMATS
415    }
416    fn default_parse_chain(&self) -> Vec<&'static str> {
417        vec![
418            "subtitle.read.srt",
419            "subtitle.read.vtt",
420            "subtitle.read.ass",
421        ]
422    }
423    fn default_process_chain(&self) -> Vec<&'static str> {
424        vec![]
425    }
426    fn assemble(&self, raw: &RawDoc) -> Document {
427        let pages = raw
428            .pages
429            .iter()
430            .map(|rp| {
431                let blocks: Vec<Block> = rp
432                    .items
433                    .iter()
434                    .map(|item| Block {
435                        kind: BlockKind::Paragraph,
436                        content: item.text.clone(),
437                        ..Default::default()
438                    })
439                    .collect();
440                let text = rp
441                    .items
442                    .iter()
443                    .map(|i| i.text.as_str())
444                    .collect::<Vec<_>>()
445                    .join("\n\n");
446                Page {
447                    name: rp.name.clone(),
448                    blocks,
449                    text,
450                    ..Default::default()
451                }
452            })
453            .collect();
454        Document {
455            pages,
456            ..Default::default()
457        }
458    }
459}
460
461/// 注册字幕后端与插件（无第三方依赖，始终编译）。
462pub fn register_defaults(reg: &mut Registry) {
463    for (id, kind) in [
464        ("subtitle.read.srt", "srt"),
465        ("subtitle.read.vtt", "vtt"),
466        ("subtitle.read.ass", "ass"),
467    ] {
468        let _ = reg.register_parse(Arc::new(SubtitleParseCapability { id, kind }), None);
469    }
470    reg.register_plugin(Arc::new(SubtitlePlugin));
471}
472
473#[cfg(test)]
474mod tests {
475    use super::*;
476
477    #[test]
478    fn clock_parses_and_normalizes() {
479        assert_eq!(parse_clock("00:00:01,000"), Some(1000));
480        assert_eq!(parse_clock("0:00:03.50"), Some(3500));
481        assert_eq!(parse_clock("01:02:03.004"), Some(3_723_004));
482        assert_eq!(parse_clock(" 1:00 "), None); // 缺小数
483        assert_eq!(parse_clock("abc"), None);
484        assert_eq!(format_clock(3_723_004), "01:02:03,004");
485        assert_eq!(format_clock(500), "00:00:00,500");
486    }
487
488    #[test]
489    fn timeline_ignores_trailing_settings() {
490        assert_eq!(
491            parse_timeline("00:00:01,000 --> 00:00:02,500 X1:0 X2:1"),
492            Some((1000, 2500))
493        );
494        assert_eq!(
495            parse_timeline("00:01.000 --> 00:02.000 line:0 align:start"),
496            Some((1000, 2000))
497        );
498        assert_eq!(parse_timeline("no arrow"), None);
499    }
500
501    #[test]
502    fn srt_cues_keep_multiline_and_strip_tags() {
503        let text = "1\n00:00:01,000 --> 00:00:04,500\n第一句\n<i>第二行</i>\n\n2\n00:00:05,000 --> 00:00:07,000\nA &amp; B\n";
504        let cues = parse_srt(text).unwrap();
505        assert_eq!(cues.len(), 2);
506        assert_eq!(cues[0].0, "00:00:01,000 --> 00:00:04,500");
507        assert_eq!(clean_text(&cues[0].1), "第一句\n第二行");
508        assert_eq!(clean_text(&cues[1].1), "A & B");
509    }
510
511    #[test]
512    fn srt_missing_arrow_blocks_are_skipped() {
513        let text = "1\n00:00:01,000 --> 00:00:02,000\n正文\n\n游离文本\n";
514        assert_eq!(parse_srt(text).unwrap().len(), 1);
515    }
516
517    #[test]
518    fn vtt_skips_note_and_requires_header() {
519        let text = "WEBVTT - 标题\n\nNOTE 注释\n不该出现\n\nc1\n00:00:01.000 --> 00:00:02.500 line:0\n<v 旁白>你好</v>\n";
520        let cues = parse_vtt(text).unwrap();
521        assert_eq!(cues.len(), 1);
522        assert_eq!(clean_text(&cues[0].1), "你好");
523
524        let err = parse_vtt("1\n00:00:01,000 --> 00:00:02,000\nx\n").unwrap_err();
525        assert_eq!(err.code.as_str(), "corrupt");
526    }
527
528    #[test]
529    fn ass_keeps_commas_in_text_and_drops_overrides() {
530        let text = "[Events]\nFormat: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text\nDialogue: 0,0:00:01.00,0:00:03.50,Default,,0,0,0,,{\\pos(1,2)}你好，世界\\N第二行\nComment: 0,0:00:01.00,0:00:02.00,Default,,0,0,0,,忽略\n";
531        let cues = parse_ass(text).unwrap();
532        assert_eq!(cues.len(), 1);
533        assert_eq!(cues[0].0, "00:00:01,000 --> 00:00:03,500");
534        assert_eq!(clean_text(&cues[0].1), "你好，世界\n第二行");
535    }
536
537    #[test]
538    fn ass_without_events_is_corrupt() {
539        let err = parse_ass("[Script Info]\nTitle: t\n").unwrap_err();
540        assert_eq!(err.code.as_str(), "corrupt");
541    }
542
543    #[test]
544    fn registry_resolves_every_subtitle_format() {
545        let mut reg = Registry::new();
546        crate::plugins::register_defaults(&mut reg);
547        let pol = crate::policy::BackendPolicy::default();
548        let default = [
549            "subtitle.read.srt",
550            "subtitle.read.vtt",
551            "subtitle.read.ass",
552        ];
553        for (fmt, want) in [
554            ("srt", "subtitle.read.srt"),
555            ("vtt", "subtitle.read.vtt"),
556            ("ass", "subtitle.read.ass"),
557            ("ssa", "subtitle.read.ass"),
558        ] {
559            let chain = reg.resolve_parse_chain(fmt, &pol, &default);
560            assert_eq!(chain.len(), 1, "{fmt} 解析链为空");
561            assert_eq!(chain[0].id(), want, "{fmt}");
562        }
563    }
564
565    #[test]
566    fn engine_parses_subtitle_end_to_end() {
567        let dir = std::env::temp_dir();
568        let cases: [(&str, &str, &str); 3] = [
569            (
570                "e2e.srt",
571                "1\n00:00:01,000 --> 00:00:02,000\n第一句\n",
572                "第一句",
573            ),
574            (
575                "e2e.vtt",
576                "WEBVTT\n\n00:00:01.000 --> 00:00:02.500\n你好\n",
577                "你好",
578            ),
579            (
580                "e2e.ass",
581                "[Events]\nFormat: Layer, Start, End, Text\nDialogue: 0,0:00:01.00,0:00:02.50,台词\n",
582                "台词",
583            ),
584        ];
585        for (name, body, want) in cases {
586            let p = dir.join(format!("docparser_{name}"));
587            std::fs::write(&p, body).unwrap();
588            let mut reg = Registry::new();
589            crate::plugins::register_defaults(&mut reg);
590            let doc = crate::engine::DocumentEngine::new(reg)
591                .parse(p.to_str().unwrap())
592                .unwrap_or_else(|e| panic!("{name} 解析失败: {e}"));
593            assert_eq!(doc.pages.len(), 1, "{name}");
594            assert_eq!(doc.pages[0].text, want, "{name}");
595            std::fs::remove_file(&p).ok();
596        }
597    }
598
599    #[test]
600    fn plugin_makes_one_page_per_cue() {
601        let raw = RawDoc {
602            pages: vec![RawPage {
603                name: Some("00:00:01,000 --> 00:00:02,000".into()),
604                items: vec![RawTextItem {
605                    text: "台词".into(),
606                    ..Default::default()
607                }],
608                ..Default::default()
609            }],
610            ..Default::default()
611        };
612        let doc = SubtitlePlugin.assemble(&raw);
613        assert_eq!(doc.pages.len(), 1);
614        assert_eq!(
615            doc.pages[0].name.as_deref(),
616            Some("00:00:01,000 --> 00:00:02,000")
617        );
618        assert_eq!(doc.pages[0].text, "台词");
619        assert_eq!(doc.pages[0].blocks.len(), 1);
620    }
621}
docparser_core/plugins/subtitle.rs

docparser_core/plugins/
subtitle.rs
