use super::types::{LineBreakMode, PlainTextConfig, PlainTextResult};
use crate::parser::document::PdfDocument;
use crate::parser::ParseResult;
use crate::text::{ExtractionOptions, TextExtractor};
use std::io::{Read, Seek};
pub struct PlainTextExtractor {
config: PlainTextConfig,
}
impl Default for PlainTextExtractor {
fn default() -> Self {
Self::new()
}
}
impl PlainTextExtractor {
pub fn new() -> Self {
Self {
config: PlainTextConfig::default(),
}
}
pub fn with_config(config: PlainTextConfig) -> Self {
Self { config }
}
pub fn extract<R: Read + Seek>(
&mut self,
document: &PdfDocument<R>,
page_index: u32,
) -> ParseResult<PlainTextResult> {
let options = ExtractionOptions {
preserve_layout: false,
space_threshold: self.config.space_threshold,
tj_space_threshold: self.config.tj_space_threshold,
newline_threshold: self.config.newline_threshold,
sort_by_position: self.config.preserve_layout,
detect_columns: false,
merge_hyphenated: matches!(
self.config.line_break_mode,
LineBreakMode::Auto | LineBreakMode::Normalize
),
..Default::default()
};
let extracted = TextExtractor::with_options(options)
.with_reading_order(self.config.preserve_layout)
.with_backward_line_wraps(self.config.preserve_layout)
.extract_from_page(document, page_index)?;
Ok(PlainTextResult::new(
self.apply_line_break_mode(&extracted.text),
))
}
pub fn extract_lines<R: Read + Seek>(
&mut self,
document: &PdfDocument<R>,
page_index: u32,
) -> ParseResult<Vec<String>> {
let result = self.extract(document, page_index)?;
Ok(result.text.lines().map(|line| line.to_string()).collect())
}
fn apply_line_break_mode(&self, text: &str) -> String {
match self.config.line_break_mode {
LineBreakMode::Auto => self.auto_line_breaks(text),
LineBreakMode::PreserveAll => text.to_string(),
LineBreakMode::Normalize => self.normalize_line_breaks(text),
}
}
fn auto_line_breaks(&self, text: &str) -> String {
let lines: Vec<&str> = text.lines().collect();
let mut result = String::with_capacity(text.len());
for (i, line) in lines.iter().enumerate() {
let trimmed = line.trim_end();
if trimmed.is_empty() {
result.push('\n');
continue;
}
result.push_str(line);
if i < lines.len() - 1 {
let next_line = lines[i + 1].trim_start();
let ends_with_punct = trimmed.ends_with('.')
|| trimmed.ends_with('!')
|| trimmed.ends_with('?')
|| trimmed.ends_with(':');
let next_is_empty = next_line.is_empty();
if ends_with_punct || next_is_empty {
result.push('\n');
} else {
result.push(' ');
}
}
}
result
}
fn normalize_line_breaks(&self, text: &str) -> String {
let lines: Vec<&str> = text.lines().collect();
let mut result = String::with_capacity(text.len());
for (i, line) in lines.iter().enumerate() {
let trimmed = line.trim_end();
if trimmed.is_empty() {
result.push('\n');
continue;
}
if trimmed.ends_with('-') && i < lines.len() - 1 {
let next_line = lines[i + 1].trim_start();
if !next_line.is_empty() {
match crate::text::extraction::hyphen_fusion_action(trimmed, next_line) {
crate::text::extraction::HyphenFusionAction::DropHyphen => {
result.push_str(&trimmed[..trimmed.len() - 1]);
continue;
}
crate::text::extraction::HyphenFusionAction::KeepHyphen => {
result.push_str(trimmed);
continue;
}
crate::text::extraction::HyphenFusionAction::NoFusion => {}
}
}
}
result.push_str(line);
if i < lines.len() - 1 {
result.push('\n');
}
}
result
}
pub fn config(&self) -> &PlainTextConfig {
&self.config
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new() {
let extractor = PlainTextExtractor::new();
assert_eq!(extractor.config.space_threshold, 0.3);
}
#[test]
fn test_with_config() {
let config = PlainTextConfig::dense();
let extractor = PlainTextExtractor::with_config(config.clone());
assert_eq!(extractor.config, config);
}
#[test]
fn test_default() {
let extractor = PlainTextExtractor::default();
assert_eq!(extractor.config, PlainTextConfig::default());
}
#[test]
fn test_normalize_line_breaks_hyphenated() {
let extractor = PlainTextExtractor::new();
let text = "This is a docu-\nment with hyphen-\nated words.";
let normalized = extractor.normalize_line_breaks(text);
assert_eq!(normalized, "This is a document with hyphenated words.");
}
#[test]
fn test_normalize_line_breaks_no_hyphen() {
let extractor = PlainTextExtractor::new();
let text = "This is a normal\ntext without\nhyphens.";
let normalized = extractor.normalize_line_breaks(text);
assert_eq!(normalized, "This is a normal\ntext without\nhyphens.");
}
#[test]
fn test_auto_line_breaks_punctuation() {
let extractor = PlainTextExtractor::new();
let text = "First sentence.\nSecond sentence.\nThird sentence.";
let processed = extractor.auto_line_breaks(text);
assert_eq!(
processed,
"First sentence.\nSecond sentence.\nThird sentence."
);
}
#[test]
fn test_auto_line_breaks_wrapped() {
let extractor = PlainTextExtractor::new();
let text = "This is a long line that\nwas wrapped in the PDF\nfor layout purposes";
let processed = extractor.auto_line_breaks(text);
assert!(processed.contains("long line that was"));
assert!(processed.contains("wrapped in the PDF for"));
}
#[test]
fn test_auto_line_breaks_empty_lines() {
let extractor = PlainTextExtractor::new();
let text = "Paragraph one.\n\nParagraph two.\n\nParagraph three.";
let processed = extractor.auto_line_breaks(text);
assert!(processed.contains("\n\n"));
}
#[test]
fn test_apply_line_break_mode_preserve_all() {
let extractor = PlainTextExtractor::with_config(PlainTextConfig {
line_break_mode: LineBreakMode::PreserveAll,
..Default::default()
});
let text = "Line 1\nLine 2\nLine 3";
let processed = extractor.apply_line_break_mode(text);
assert_eq!(processed, text);
}
#[test]
fn test_apply_line_break_mode_normalize() {
let extractor = PlainTextExtractor::with_config(PlainTextConfig {
line_break_mode: LineBreakMode::Normalize,
..Default::default()
});
let text = "docu-\nment";
let processed = extractor.apply_line_break_mode(text);
assert_eq!(processed, "document");
}
#[test]
fn test_apply_line_break_mode_auto() {
let extractor = PlainTextExtractor::with_config(PlainTextConfig {
line_break_mode: LineBreakMode::Auto,
..Default::default()
});
let text = "First sentence.\nSecond part";
let processed = extractor.apply_line_break_mode(text);
assert!(processed.contains("First sentence.\nSecond"));
}
#[test]
fn test_config_getter() {
let config = PlainTextConfig::loose();
let extractor = PlainTextExtractor::with_config(config.clone());
assert_eq!(extractor.config(), &config);
}
}