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
12// ================================================================
13// Implementing Loadable trait for loading epubs
14// ================================================================
15
16loadable_trait!(
17    Loadable,
18    EpubLoaderError,
19    EpubDoc<BufReader<File>>,
20    load,
21    load_with_path
22);
23
24impl Loadable for PathBuf {
25    fn load(self) -> Result<EpubDoc<BufReader<File>>, EpubLoaderError> {
26        EpubDoc::new(self).map_err(EpubLoaderError::EpubError)
27    }
28
29    fn load_with_path(self) -> Result<(PathBuf, EpubDoc<BufReader<File>>), EpubLoaderError> {
30        let contents = EpubDoc::new(&self).map_err(EpubLoaderError::EpubError);
31        Ok((self, contents?))
32    }
33}
34
35// ================================================================
36// EpubFileLoader definitions and implementations
37// ================================================================
38
39/// [EpubFileLoader] is a utility for loading epub files from the filesystem using glob patterns or
40///  directory paths. It provides methods to read file contents and handle errors gracefully.
41///
42/// # Errors
43///
44/// This module defines a custom error type [EpubLoaderError] which can represent various errors
45///  that might occur during file loading operations, such as any [FileLoaderError](crate::loaders::file::FileLoaderError) alongside
46///  specific EPUB-related errors.
47///
48/// # Example Usage
49///
50/// ```no_run
51/// use rig_core::loaders::{EpubFileLoader, RawTextProcessor, StripXmlProcessor};
52///
53/// fn main() -> Result<(), Box<dyn std::error::Error>> {
54///     // Create a FileLoader using a glob pattern
55///     let loader = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?;
56///
57///     // Load epub file contents by chapter, ignoring any errors
58///     let contents = loader
59///         .load_with_path()
60///         .ignore_errors()
61///         .by_chapter()
62///         .ignore_errors();
63///
64///     for (path, chapters) in contents {
65///         println!("{}", path.display());
66///         for (idx, chapter) in chapters {
67///             println!("Chapter {} begins", idx);
68///             println!("{}", chapter);
69///             println!("Chapter {} ends", idx);
70///         }
71///     }
72///
73///     // Create a FileLoader using a glob pattern with stripping xml
74///     let loader = EpubFileLoader::<_, StripXmlProcessor>::with_glob("tests/data/*.epub")?;
75///
76///     // Load epub file contents by chapter, ignoring any errors
77///     let contents = loader
78///         .load_with_path()
79///         .ignore_errors()
80///         .by_chapter()
81///         .ignore_errors();
82///
83///     for (path, chapters) in contents {
84///         println!("{}", path.display());
85///         for (idx, chapter) in chapters {
86///             println!("Chapter {} begins", idx);
87///             println!("{}", chapter);
88///             println!("Chapter {} ends", idx);
89///         }
90///     }
91///
92///     Ok(())
93/// }
94/// ```
95///
96/// [EpubFileLoader] uses strict typing between the iterator methods to ensure that transitions
97///  between different implementations of the loaders and it's methods are handled properly by
98///  the compiler.
99pub struct EpubFileLoader<'a, T, P = RawTextProcessor> {
100    iterator: Box<dyn Iterator<Item = T> + 'a>,
101    _processor: PhantomData<P>,
102}
103
104type EpubLoaded = Result<(PathBuf, EpubDoc<BufReader<File>>), EpubLoaderError>;
105
106impl<'a, P> EpubFileLoader<'a, Result<PathBuf, EpubLoaderError>, P> {
107    /// Loads the contents of the epub files within the iterator returned by [EpubFileLoader::with_glob]
108    ///  or [EpubFileLoader::with_dir]. Loaded EPUB documents are raw EPUB instances that can be
109    ///  further processed (by chapter, etc).
110    ///
111    /// # Example
112    /// Load epub files in directory "tests/data/*.epub" and return the loaded documents
113    ///
114    /// ```no_run
115    /// use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
116    ///
117    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
118    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?.load().into_iter();
119    /// for result in content {
120    ///     match result {
121    ///         Ok(doc) => println!("{:?}", doc),
122    ///         Err(e) => eprintln!("Error reading epub: {}", e),
123    ///     }
124    /// }
125    /// # Ok(())
126    /// # }
127    /// ```
128    pub fn load(self) -> EpubFileLoader<'a, Result<EpubDoc<BufReader<File>>, EpubLoaderError>, P> {
129        EpubFileLoader {
130            iterator: Box::new(self.iterator.map(|res| res.load())),
131            _processor: PhantomData,
132        }
133    }
134
135    /// Loads the contents of the epub files within the iterator returned by [EpubFileLoader::with_glob]
136    ///  or [EpubFileLoader::with_dir]. Loaded EPUB documents are raw EPUB instances with their path
137    ///  that can be further processed.
138    ///
139    /// # Example
140    /// Load epub files in directory "tests/data/*.epub" and return the loaded documents
141    ///
142    /// ```no_run
143    /// use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
144    ///
145    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
146    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?.load_with_path().into_iter();
147    /// for result in content {
148    ///     match result {
149    ///         Ok((path, doc)) => println!("{:?} {:?}", path, doc),
150    ///         Err(e) => eprintln!("Error reading epub: {}", e),
151    ///     }
152    /// }
153    /// # Ok(())
154    /// # }
155    /// ```
156    pub fn load_with_path(self) -> EpubFileLoader<'a, EpubLoaded, P> {
157        EpubFileLoader {
158            iterator: Box::new(self.iterator.map(|res| res.load_with_path())),
159            _processor: PhantomData,
160        }
161    }
162}
163
164impl<'a, P> EpubFileLoader<'a, Result<PathBuf, EpubLoaderError>, P>
165where
166    P: TextProcessor,
167{
168    /// Directly reads the contents of the epub files within the iterator returned by
169    ///  [EpubFileLoader::with_glob] or [EpubFileLoader::with_dir].
170    ///
171    /// # Example
172    /// Read epub files in directory "tests/data/*.epub" and return the contents of the documents.
173    ///
174    /// ```no_run
175    /// # use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
176    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
177    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?.read().into_iter();
178    /// for result in content {
179    ///     match result {
180    ///         Ok(content) => println!("{}", content),
181    ///         Err(e) => eprintln!("Error reading epub: {}", e),
182    ///     }
183    /// }
184    /// # Ok(())
185    /// # }
186    /// ```
187    pub fn read(self) -> EpubFileLoader<'a, Result<String, EpubLoaderError>, P> {
188        EpubFileLoader {
189            iterator: Box::new(self.iterator.map(|res| {
190                let doc = res.load().map(EpubChapterIterator::<P>::from)?;
191
192                Ok(doc
193                    .into_iter()
194                    .collect::<Result<Vec<String>, EpubLoaderError>>()?
195                    .into_iter()
196                    .collect::<String>())
197            })),
198            _processor: PhantomData,
199        }
200    }
201
202    /// Directly reads the contents of the epub files within the iterator returned by
203    ///  [EpubFileLoader::with_glob] or [EpubFileLoader::with_dir] and returns the path along with
204    ///  the content.
205    ///
206    /// # Example
207    /// Read epub files in directory "tests/data/*.epub" and return the content and paths of the documents.
208    ///
209    /// ```no_run
210    /// # use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
211    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
212    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?.read_with_path().into_iter();
213    /// for result in content {
214    ///     match result {
215    ///         Ok((path, content)) => println!("{:?} {}", path, content),
216    ///         Err(e) => eprintln!("Error reading epub: {}", e),
217    ///     }
218    /// }
219    /// # Ok(())
220    /// # }
221    /// ```
222    pub fn read_with_path(
223        self,
224    ) -> EpubFileLoader<'a, Result<(PathBuf, String), EpubLoaderError>, P> {
225        EpubFileLoader {
226            iterator: Box::new(self.iterator.map(|res| {
227                let (path, doc) = res.load_with_path()?;
228
229                let content = EpubChapterIterator::<P>::from(doc)
230                    .collect::<Result<Vec<String>, EpubLoaderError>>()?
231                    .into_iter()
232                    .collect::<String>();
233                Ok((path, content))
234            })),
235            _processor: PhantomData,
236        }
237    }
238}
239
240impl<'a, P> EpubFileLoader<'a, EpubDoc<BufReader<File>>, P>
241where
242    P: TextProcessor + 'a,
243{
244    /// Chunks the chapters of a loaded document by chapter, flattened as a single vector.
245    ///
246    /// # Example
247    /// Load epub files in directory "tests/data/*.epub" and chunk all document into it's chapters.
248    ///
249    /// ```no_run
250    /// # use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
251    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
252    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?
253    ///     .load()
254    ///     .ignore_errors()
255    ///     .by_chapter()
256    ///     .into_iter();
257    /// for result in content {
258    ///     match result {
259    ///         Ok(chapter) => println!("{}", chapter),
260    ///         Err(e) => eprintln!("Error reading chapter: {}", e),
261    ///     }
262    /// }
263    /// # Ok(())
264    /// # }
265    /// ```
266    pub fn by_chapter(self) -> EpubFileLoader<'a, Result<String, EpubLoaderError>, P> {
267        EpubFileLoader {
268            iterator: Box::new(self.iterator.flat_map(EpubChapterIterator::<P>::from)),
269            _processor: PhantomData,
270        }
271    }
272}
273
274type ByChapter = (PathBuf, Vec<(usize, Result<String, EpubLoaderError>)>);
275impl<'a, P: TextProcessor> EpubFileLoader<'a, (PathBuf, EpubDoc<BufReader<File>>), P> {
276    /// Chunks the chapters of a loaded document by chapter, processed as a vector of documents by path
277    ///  which each document container an inner vector of chapters by chapter number.
278    ///
279    /// # Example
280    /// Read epub files in directory "tests/data/*.epub" and chunk all documents by path by it's chapters.
281    ///
282    /// ```no_run
283    /// # use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
284    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
285    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?
286    ///     .load_with_path()
287    ///     .ignore_errors()
288    ///     .by_chapter()
289    ///     .ignore_errors()
290    ///     .into_iter();
291    ///
292    /// for result in content {
293    ///     println!("{:?}", result);
294    /// }
295    /// # Ok(())
296    /// # }
297    /// ```
298    pub fn by_chapter(self) -> EpubFileLoader<'a, ByChapter, P> {
299        EpubFileLoader {
300            iterator: Box::new(self.iterator.map(|doc| {
301                let (path, doc) = doc;
302
303                (
304                    path,
305                    EpubChapterIterator::<P>::from(doc)
306                        .enumerate()
307                        .collect::<Vec<_>>(),
308                )
309            })),
310            _processor: PhantomData,
311        }
312    }
313}
314
315impl<'a, P> EpubFileLoader<'a, ByChapter, P>
316where
317    P: TextProcessor,
318{
319    /// Ignores errors in the iterator, returning only successful results. This can be used on any
320    ///  [EpubFileLoader] state of iterator whose items are results.
321    ///
322    /// # Example
323    /// Read files in directory "tests/data/*.epub" and ignore errors from unreadable files.
324    ///
325    /// ```no_run
326    /// # use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
327    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
328    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?
329    ///     .load_with_path()
330    ///     .ignore_errors()
331    ///     .by_chapter()
332    ///     .ignore_errors();
333    /// for (_path, chapters) in content {
334    ///     println!("{}", chapters.len())
335    /// }
336    /// # Ok(())
337    /// # }
338    /// ```
339    pub fn ignore_errors(self) -> EpubFileLoader<'a, (PathBuf, Vec<(usize, String)>), P> {
340        EpubFileLoader {
341            iterator: Box::new(self.iterator.map(|(path, chapters)| {
342                let chapters = chapters
343                    .into_iter()
344                    .filter_map(|(idx, res)| res.ok().map(|content| (idx, content)))
345                    .collect::<Vec<_>>();
346                (path, chapters)
347            })),
348            _processor: PhantomData,
349        }
350    }
351}
352
353loader_scaffold!(EpubFileLoader, EpubLoaderError, dir: all_entries, extra: P);
354
355// ================================================================
356// EpubChapterIterator definitions and implementations
357// ================================================================
358
359struct EpubChapterIterator<P> {
360    epub: EpubDoc<BufReader<File>>,
361    finished: bool,
362    _processor: PhantomData<P>,
363}
364
365impl<P> From<EpubDoc<BufReader<File>>> for EpubChapterIterator<P> {
366    fn from(epub: EpubDoc<BufReader<File>>) -> Self {
367        Self::new(epub)
368    }
369}
370
371impl<P> EpubChapterIterator<P> {
372    fn new(epub: EpubDoc<BufReader<File>>) -> Self {
373        Self {
374            epub,
375            finished: false,
376            _processor: PhantomData,
377        }
378    }
379}
380
381impl<P> Iterator for EpubChapterIterator<P>
382where
383    P: TextProcessor,
384{
385    type Item = Result<String, EpubLoaderError>;
386
387    fn next(&mut self) -> Option<Self::Item> {
388        if self.finished {
389            return None;
390        }
391
392        // ignore empty chapters if they exist
393        while !self.finished {
394            let chapter = self.epub.get_current_str();
395
396            if !self.epub.go_next() {
397                self.finished = true;
398            }
399
400            if let Some((text, _)) = chapter {
401                return Some(
402                    P::process(&text)
403                        .map_err(|err| EpubLoaderError::TextProcessorError(Box::new(err))),
404                );
405            }
406        }
407
408        None
409    }
410}
411
412#[cfg(test)]
413mod tests {
414    use crate::loaders::epub::RawTextProcessor;
415    use crate::loaders::test_fixtures::{fixture_glob, fixture_path};
416
417    use super::EpubFileLoader;
418
419    #[test]
420    fn test_epub_loader_with_errors() {
421        let glob = fixture_glob("*.epub");
422        let loader = EpubFileLoader::<_, RawTextProcessor>::with_glob(&glob).unwrap();
423        let actual = loader
424            .load_with_path()
425            .ignore_errors()
426            .by_chapter()
427            .into_iter()
428            .collect::<Vec<_>>();
429
430        assert_eq!(actual.len(), 1);
431
432        let (_, chapters) = &actual[0];
433        assert_eq!(chapters.len(), 3);
434
435        for chapter in chapters {
436            assert!(chapter.1.is_ok());
437        }
438    }
439
440    #[test]
441    fn test_epub_loader_with_ignoring_errors() {
442        let glob = fixture_glob("*.epub");
443        let loader = EpubFileLoader::<_, RawTextProcessor>::with_glob(&glob).unwrap();
444        let actual = loader
445            .load_with_path()
446            .ignore_errors()
447            .by_chapter()
448            .ignore_errors()
449            .into_iter()
450            .collect::<Vec<_>>();
451
452        assert_eq!(actual.len(), 1);
453
454        let (_, chapters) = &actual[0];
455        assert_eq!(chapters.len(), 3);
456    }
457
458    #[test]
459    fn test_single_file() {
460        let glob = fixture_glob("*.epub");
461        let loader = EpubFileLoader::<_, RawTextProcessor>::with_glob(&glob).unwrap();
462
463        let actual = loader
464            .read()
465            .ignore_errors()
466            .into_iter()
467            .collect::<Vec<_>>();
468
469        assert_eq!(actual.len(), 1);
470    }
471
472    #[test]
473    fn test_single_file_with_path() {
474        let glob = fixture_glob("*.epub");
475        let loader = EpubFileLoader::<_, RawTextProcessor>::with_glob(&glob).unwrap();
476
477        let actual = loader
478            .read_with_path()
479            .ignore_errors()
480            .into_iter()
481            .collect::<Vec<_>>();
482
483        assert_eq!(actual.len(), 1);
484
485        let (path, _) = &actual[0];
486        assert_eq!(path, &fixture_path("dummy.epub"));
487    }
488}
rig_core/loaders/epub/loader.rs

rig_core/loaders/epub/
loader.rs
