Skip to main content

kobo_core/html_text/
extract.rs

1//! DOM extraction: chapter XHTML -> plain text + per-block-element segments.
2
3use scraper::{Html, Selector};
4use std::sync::LazyLock;
5
6/// A run of chapter text belonging to one block element, with its char range.
7/// For `<img>`/`<figure>` segments the text range is empty (zero-width marker at
8/// the image's flow position) and `src`/`caption` carry the image data.
9#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
10pub struct TextSegment {
11    /// Inclusive start byte/char offset into the chapter text.
12    pub start: usize,
13    /// Exclusive end offset.
14    pub end: usize,
15    /// Element tag (e.g. `p`, `h1`, `li`, `img`, `figure`).
16    pub tag: String,
17    /// The element's `id` attribute if any (useful for anchoring).
18    pub id: Option<String>,
19    /// `src` for `img`/`figure` (None for text blocks).
20    pub src: Option<String>,
21    /// `figcaption` text for a `figure` (None for non-figures).
22    pub caption: Option<String>,
23}
24
25impl TextSegment {
26    /// Does this segment contain char offset `i`?
27    pub fn contains(&self, i: usize) -> bool {
28        i >= self.start && i < self.end
29    }
30}
31
32/// Selector for the block elements we extract as flow items (text + images).
33/// `image` matches SVG `<image>` (covers, inline SVG art).
34const BLOCK_SELECTOR: &str = "p, h1, h2, h3, h4, h5, h6, li, blockquote, div, figure, img, image";
35const CONTAINER_TAGS: &[&str] = &["div", "blockquote"];
36
37static BLOCK_SELECTOR_OBJ: LazyLock<Selector> =
38    LazyLock::new(|| Selector::parse(BLOCK_SELECTOR).expect("valid selector"));
39static IMG_SELECTOR_OBJ: LazyLock<Selector> =
40    LazyLock::new(|| Selector::parse("img").expect("selector"));
41static CAP_SELECTOR_OBJ: LazyLock<Selector> =
42    LazyLock::new(|| Selector::parse("figcaption").expect("selector"));
43
44/// Extract plain text + per-block-element segments from chapter XHTML.
45/// Images (`<img>`/`<figure>`) appear in document order as zero-width segments
46/// carrying `src`/`caption` (they contribute no text to the returned string).
47pub fn extract(xhtml: &str) -> (String, Vec<TextSegment>) {
48    let html = Html::parse_fragment(xhtml);
49    let sel = BLOCK_SELECTOR_OBJ.clone();
50    let child_sel = BLOCK_SELECTOR_OBJ.clone();
51    let img_sel = IMG_SELECTOR_OBJ.clone();
52    let cap_sel = CAP_SELECTOR_OBJ.clone();
53
54    let mut text = String::new();
55    let mut segs: Vec<TextSegment> = Vec::new();
56    let mut figure_srcs: std::collections::HashSet<String> = std::collections::HashSet::new();
57
58    for elem in html.select(&sel) {
59        let tag = elem.value().name().to_string();
60        let pos = text.len();
61        let id = elem.value().attr("id").map(|s| s.to_string());
62
63        let seg = match tag.as_str() {
64            "figure" => extract_figure_segment(elem, &img_sel, &cap_sel, pos, id, &mut figure_srcs),
65            "img" => {
66                if figure_srcs.contains(elem.value().attr("src").unwrap_or("")) {
67                    None
68                } else {
69                    Some(TextSegment {
70                        start: pos,
71                        end: pos,
72                        tag,
73                        id,
74                        src: elem.value().attr("src").map(|s| s.to_string()),
75                        caption: None,
76                    })
77                }
78            }
79            "image" => extract_svg_image_segment(elem, pos, id),
80            _ => extract_text_segment(elem, &child_sel, &mut text, pos, tag, id, &segs),
81        };
82        if let Some(s) = seg {
83            segs.push(s);
84        }
85    }
86
87    let extra = scan_raw_svg_images(xhtml, &segs, &mut text);
88    segs.extend(extra);
89
90    (text, segs)
91}
92
93fn extract_figure_segment(
94    elem: scraper::ElementRef,
95    img_sel: &Selector,
96    cap_sel: &Selector,
97    pos: usize,
98    id: Option<String>,
99    figure_srcs: &mut std::collections::HashSet<String>,
100) -> Option<TextSegment> {
101    let src = elem
102        .select(img_sel)
103        .next()
104        .and_then(|i| i.value().attr("src"))
105        .map(|s| {
106            figure_srcs.insert(s.to_string());
107            s.to_string()
108        });
109    let cap = elem
110        .select(cap_sel)
111        .next()
112        .map(|c| c.text().collect::<String>().trim().to_string());
113    Some(TextSegment {
114        start: pos,
115        end: pos,
116        tag: "figure".to_string(),
117        id,
118        src,
119        caption: cap,
120    })
121}
122
123fn extract_svg_image_segment(
124    elem: scraper::ElementRef,
125    pos: usize,
126    id: Option<String>,
127) -> Option<TextSegment> {
128    let src = elem
129        .value()
130        .attr("xlink:href")
131        .or_else(|| elem.value().attr("href"))
132        .unwrap_or("");
133    if src.is_empty() {
134        return None;
135    }
136    Some(TextSegment {
137        start: pos,
138        end: pos,
139        tag: "image".to_string(),
140        id,
141        src: Some(src.to_string()),
142        caption: None,
143    })
144}
145
146fn extract_text_segment(
147    elem: scraper::ElementRef,
148    child_sel: &Selector,
149    text: &mut String,
150    pos: usize,
151    tag: String,
152    id: Option<String>,
153    segs: &[TextSegment],
154) -> Option<TextSegment> {
155    if CONTAINER_TAGS.contains(&tag.as_str()) && elem.select(child_sel).next().is_some() {
156        return None;
157    }
158    let raw: String = elem.text().collect();
159    let t = raw.trim();
160    if t.is_empty() {
161        return None;
162    }
163    if !text.is_empty() {
164        text.push('\n');
165    }
166    let content_start = text.len();
167    text.push_str(t);
168    let end = text.len();
169    let dup = segs
170        .last()
171        .map(|s| text[s.start..s.end] == text[content_start..end])
172        .unwrap_or(false);
173    if dup {
174        text.truncate(pos);
175        return None;
176    }
177    Some(TextSegment {
178        start: content_start,
179        end,
180        tag,
181        id,
182        src: None,
183        caption: None,
184    })
185}
186
187fn scan_raw_svg_images(xhtml: &str, segs: &[TextSegment], text: &mut String) -> Vec<TextSegment> {
188    let captured_srcs: std::collections::HashSet<String> =
189        segs.iter().filter_map(|s| s.src.clone()).collect();
190    let mut extra: Vec<TextSegment> = Vec::new();
191    let mut search = 0;
192    while let Some(pos) = xhtml[search..].find("<image") {
193        let abs = search + pos;
194        let tag_end = xhtml[abs..]
195            .find('>')
196            .map(|p| abs + p)
197            .unwrap_or(xhtml.len());
198        let tag = &xhtml[abs..tag_end];
199        for attr in &["xlink:href", "href"] {
200            if let Some(eq) = tag.find(attr) {
201                let rest = &tag[eq + attr.len()..];
202                let rest = rest.trim_start();
203                if let Some(after_eq) = rest.strip_prefix('=') {
204                    let v = after_eq.trim_start();
205                    let q = v.chars().next().unwrap_or('"');
206                    if (q == '"' || q == '\'') && v.len() > 1 {
207                        if let Some(end) = v[1..].find(q) {
208                            let src = &v[1..1 + end];
209                            if !captured_srcs.contains(src) && !src.is_empty() {
210                                extra.push(TextSegment {
211                                    start: text.len(),
212                                    end: text.len(),
213                                    tag: "image".to_string(),
214                                    id: None,
215                                    src: Some(src.to_string()),
216                                    caption: None,
217                                });
218                            }
219                            break;
220                        }
221                    }
222                }
223            }
224        }
225        search = tag_end;
226    }
227    extra
228}
229
230/// Find the [`TextSegment`] containing char offset `i`, if any.
231pub fn segment_at(segs: &[TextSegment], i: usize) -> Option<&TextSegment> {
232    segs.iter().find(|s| s.contains(i))
233}
234#[cfg(test)]
235mod tests;