rig_core/loaders/epub/
loader.rs

1use epub::doc::EpubDoc;
2
3use std::fs::File;
4use std::io::BufReader;
5use std::marker::PhantomData;
6use std::path::PathBuf;
7
8use super::RawTextProcessor;
9use super::errors::EpubLoaderError;
10use super::text_processors::TextProcessor;
11
12loadable_trait!(
13    Loadable,
14    EpubLoaderError,
15    EpubDoc<BufReader<File>>,
16    load,
17    load_with_path
18);
19
20impl Loadable for PathBuf {
21    fn load(self) -> Result<EpubDoc<BufReader<File>>, EpubLoaderError> {
22        EpubDoc::new(self).map_err(EpubLoaderError::EpubError)
23    }
24
25    fn load_with_path(self) -> Result<(PathBuf, EpubDoc<BufReader<File>>), EpubLoaderError> {
26        let contents = EpubDoc::new(&self).map_err(EpubLoaderError::EpubError);
27        Ok((self, contents?))
28    }
29}
30
31/// Iterator pipeline for synchronous EPUB loading and chapter processing.
32/// `P` selects text processing and defaults to [`RawTextProcessor`].
33/// Loading and processing errors are yielded unless filtered; unavailable chapter
34/// text is skipped.
35pub struct EpubFileLoader<'a, T, P = RawTextProcessor> {
36    iterator: Box<dyn Iterator<Item = T> + 'a>,
37    _processor: PhantomData<P>,
38}
39
40type EpubLoaded = Result<(PathBuf, EpubDoc<BufReader<File>>), EpubLoaderError>;
41
42impl<'a, P> EpubFileLoader<'a, Result<PathBuf, EpubLoaderError>, P> {
43    /// Loads each EPUB during iteration, yielding document-loading errors.
44    ///
45    /// ```no_run
46    /// use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
47    ///
48    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
49    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?.load().into_iter();
50    /// for result in content {
51    ///     match result {
52    ///         Ok(doc) => println!("{doc:?}"),
53    ///         Err(e) => eprintln!("Error reading epub: {e}"),
54    ///     }
55    /// }
56    /// # Ok(())
57    /// # }
58    /// ```
59    pub fn load(self) -> EpubFileLoader<'a, Result<EpubDoc<BufReader<File>>, EpubLoaderError>, P> {
60        EpubFileLoader {
61            iterator: Box::new(self.iterator.map(Loadable::load)),
62            _processor: PhantomData,
63        }
64    }
65
66    /// Loads each EPUB paired with its source path, yielding loading errors.
67    ///
68    /// ```no_run
69    /// use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
70    ///
71    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
72    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?.load_with_path().into_iter();
73    /// for result in content {
74    ///     match result {
75    ///         Ok((path, doc)) => println!("{path:?} {doc:?}"),
76    ///         Err(e) => eprintln!("Error reading epub: {e}"),
77    ///     }
78    /// }
79    /// # Ok(())
80    /// # }
81    /// ```
82    pub fn load_with_path(self) -> EpubFileLoader<'a, EpubLoaded, P> {
83        EpubFileLoader {
84            iterator: Box::new(self.iterator.map(Loadable::load_with_path)),
85            _processor: PhantomData,
86        }
87    }
88}
89
90impl<'a, P> EpubFileLoader<'a, Result<PathBuf, EpubLoaderError>, P>
91where
92    P: TextProcessor,
93{
94    /// Concatenates processed chapters without separators, yielding loading or
95    /// processing errors. Unavailable chapter text is skipped.
96    ///
97    /// ```no_run
98    /// # use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
99    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
100    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?.read().into_iter();
101    /// for result in content {
102    ///     match result {
103    ///         Ok(content) => println!("{content}"),
104    ///         Err(e) => eprintln!("Error reading epub: {e}"),
105    ///     }
106    /// }
107    /// # Ok(())
108    /// # }
109    /// ```
110    pub fn read(self) -> EpubFileLoader<'a, Result<String, EpubLoaderError>, P> {
111        EpubFileLoader {
112            iterator: Box::new(self.iterator.map(|res| {
113                let doc = res.load().map(EpubChapterIterator::<P>::from)?;
114
115                Ok(doc
116                    .into_iter()
117                    .collect::<Result<Vec<String>, EpubLoaderError>>()?
118                    .into_iter()
119                    .collect::<String>())
120            })),
121            _processor: PhantomData,
122        }
123    }
124
125    /// Pairs each path with concatenated processed chapters, yielding loading
126    /// or processing errors. Unavailable chapter text is skipped.
127    ///
128    /// ```no_run
129    /// # use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
130    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
131    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?.read_with_path().into_iter();
132    /// for result in content {
133    ///     match result {
134    ///         Ok((path, content)) => println!("{path:?} {content}"),
135    ///         Err(e) => eprintln!("Error reading epub: {e}"),
136    ///     }
137    /// }
138    /// # Ok(())
139    /// # }
140    /// ```
141    pub fn read_with_path(
142        self,
143    ) -> EpubFileLoader<'a, Result<(PathBuf, String), EpubLoaderError>, P> {
144        EpubFileLoader {
145            iterator: Box::new(self.iterator.map(|res| {
146                let (path, doc) = res.load_with_path()?;
147
148                let content = EpubChapterIterator::<P>::from(doc)
149                    .collect::<Result<Vec<String>, EpubLoaderError>>()?
150                    .into_iter()
151                    .collect::<String>();
152                Ok((path, content))
153            })),
154            _processor: PhantomData,
155        }
156    }
157}
158
159impl<'a, P> EpubFileLoader<'a, EpubDoc<BufReader<File>>, P>
160where
161    P: TextProcessor + 'a,
162{
163    /// Yields processed chapter results across documents in input order,
164    /// skipping unavailable chapter text.
165    ///
166    /// ```no_run
167    /// # use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
168    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
169    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?
170    ///     .load()
171    ///     .ignore_errors()
172    ///     .by_chapter()
173    ///     .into_iter();
174    /// for result in content {
175    ///     match result {
176    ///         Ok(chapter) => println!("{chapter}"),
177    ///         Err(e) => eprintln!("Error reading chapter: {e}"),
178    ///     }
179    /// }
180    /// # Ok(())
181    /// # }
182    /// ```
183    pub fn by_chapter(self) -> EpubFileLoader<'a, Result<String, EpubLoaderError>, P> {
184        EpubFileLoader {
185            iterator: Box::new(self.iterator.flat_map(EpubChapterIterator::<P>::from)),
186            _processor: PhantomData,
187        }
188    }
189}
190
191type ByChapter = (PathBuf, Vec<(usize, Result<String, EpubLoaderError>)>);
192impl<'a, P: TextProcessor> EpubFileLoader<'a, (PathBuf, EpubDoc<BufReader<File>>), P> {
193    /// Pairs each path with zero-based indices and chapter-processing results.
194    /// Indices count available chapter texts, not skipped source entries.
195    ///
196    /// ```no_run
197    /// # use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
198    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
199    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?
200    ///     .load_with_path()
201    ///     .ignore_errors()
202    ///     .by_chapter()
203    ///     .ignore_errors()
204    ///     .into_iter();
205    ///
206    /// for result in content {
207    ///     println!("{result:?}");
208    /// }
209    /// # Ok(())
210    /// # }
211    /// ```
212    pub fn by_chapter(self) -> EpubFileLoader<'a, ByChapter, P> {
213        EpubFileLoader {
214            iterator: Box::new(self.iterator.map(|doc| {
215                let (path, doc) = doc;
216
217                (
218                    path,
219                    EpubChapterIterator::<P>::from(doc)
220                        .enumerate()
221                        .collect::<Vec<_>>(),
222                )
223            })),
224            _processor: PhantomData,
225        }
226    }
227}
228
229impl<'a, P> EpubFileLoader<'a, ByChapter, P>
230where
231    P: TextProcessor,
232{
233    /// Drops failed chapter results, retaining paths and existing chapter indices.
234    /// Documents with no successful chapters remain.
235    ///
236    /// ```no_run
237    /// # use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
238    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
239    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?
240    ///     .load_with_path()
241    ///     .ignore_errors()
242    ///     .by_chapter()
243    ///     .ignore_errors();
244    /// for (_path, chapters) in content {
245    ///     println!("{}", chapters.len())
246    /// }
247    /// # Ok(())
248    /// # }
249    /// ```
250    pub fn ignore_errors(self) -> EpubFileLoader<'a, (PathBuf, Vec<(usize, String)>), P> {
251        EpubFileLoader {
252            iterator: Box::new(self.iterator.map(|(path, chapters)| {
253                let chapters = chapters
254                    .into_iter()
255                    .filter_map(|(idx, res)| res.ok().map(|content| (idx, content)))
256                    .collect::<Vec<_>>();
257                (path, chapters)
258            })),
259            _processor: PhantomData,
260        }
261    }
262}
263
264loader_scaffold!(EpubFileLoader, EpubLoaderError, dir: all_entries, extra: P);
265
266struct EpubChapterIterator<P> {
267    epub: EpubDoc<BufReader<File>>,
268    finished: bool,
269    _processor: PhantomData<P>,
270}
271
272impl<P> From<EpubDoc<BufReader<File>>> for EpubChapterIterator<P> {
273    fn from(epub: EpubDoc<BufReader<File>>) -> Self {
274        Self::new(epub)
275    }
276}
277
278impl<P> EpubChapterIterator<P> {
279    fn new(epub: EpubDoc<BufReader<File>>) -> Self {
280        Self {
281            epub,
282            finished: false,
283            _processor: PhantomData,
284        }
285    }
286}
287
288impl<P> Iterator for EpubChapterIterator<P>
289where
290    P: TextProcessor,
291{
292    type Item = Result<String, EpubLoaderError>;
293
294    fn next(&mut self) -> Option<Self::Item> {
295        if self.finished {
296            return None;
297        }
298
299        // Unavailable chapter text is skipped rather than yielded as an error.
300        while !self.finished {
301            let chapter = self.epub.get_current_str();
302
303            if !self.epub.go_next() {
304                self.finished = true;
305            }
306
307            if let Some((text, _)) = chapter {
308                return Some(
309                    P::process(&text)
310                        .map_err(|err| EpubLoaderError::TextProcessorError(Box::new(err))),
311                );
312            }
313        }
314
315        None
316    }
317}
318
319#[cfg(test)]
320mod tests;
rig_core/loaders/epub/loader.rs

rig_core/loaders/epub/
loader.rs
