kcode_doc_extraction/
lib.rs

1//! Local PDF, DOCX, spreadsheet, and plain-text extraction.
2
3#![deny(missing_docs)]
4#![forbid(unsafe_code)]
5
6use std::{
7    error::Error as StdError,
8    fmt,
9    io::{Cursor, Read},
10};
11
12use calamine::{Reader as _, open_workbook_auto_from_rs};
13use quick_xml::{Reader as XmlReader, events::Event as XmlEvent};
14
15/// Supported local document representation.
16#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
17pub enum DocumentFormat {
18    /// Searchable PDF text extraction without OCR.
19    Pdf,
20    /// Microsoft Word Open XML document.
21    Docx,
22    /// XLSX, XLS, XLSB, or ODS workbook.
23    Spreadsheet,
24    /// UTF-8-family text, including CSV, TSV, Markdown, JSON, YAML, and XML.
25    PlainText,
26}
27
28impl DocumentFormat {
29    /// Returns the stable format label.
30    pub const fn as_str(self) -> &'static str {
31        match self {
32            Self::Pdf => "pdf",
33            Self::Docx => "docx",
34            Self::Spreadsheet => "spreadsheet",
35            Self::PlainText => "text",
36        }
37    }
38}
39
40/// Stable document-extraction failure classification.
41#[derive(Clone, Copy, Debug, Eq, PartialEq)]
42pub enum ErrorKind {
43    /// Input metadata, size, or bytes were invalid.
44    InvalidInput,
45    /// The filename and media type do not identify a supported format.
46    UnsupportedFormat,
47    /// A local parser could not extract the document.
48    ExtractionFailed,
49    /// The parser found no readable text.
50    EmptyText,
51}
52
53/// A local document validation or extraction failure.
54#[derive(Debug)]
55pub struct Error {
56    kind: ErrorKind,
57    message: String,
58}
59
60impl Error {
61    fn new(kind: ErrorKind, message: impl Into<String>) -> Self {
62        Self {
63            kind,
64            message: message.into(),
65        }
66    }
67
68    /// Returns the stable failure classification.
69    pub const fn kind(&self) -> ErrorKind {
70        self.kind
71    }
72
73    /// Returns the parser or validation detail.
74    pub fn message(&self) -> &str {
75        &self.message
76    }
77}
78
79impl fmt::Display for Error {
80    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
81        formatter.write_str(&self.message)
82    }
83}
84
85impl StdError for Error {}
86
87/// Result type returned by document extraction.
88pub type Result<T> = std::result::Result<T, Error>;
89
90/// Input and output resource limits.
91#[derive(Clone, Debug, Eq, PartialEq)]
92pub struct ExtractionLimits {
93    /// Maximum in-memory document bytes.
94    pub max_bytes: usize,
95    /// Maximum returned Unicode characters.
96    pub max_characters: usize,
97}
98
99impl Default for ExtractionLimits {
100    fn default() -> Self {
101        Self {
102            max_bytes: 20 * 1024 * 1024,
103            max_characters: 1_000_000,
104        }
105    }
106}
107
108/// One complete in-memory document.
109#[derive(Clone, Eq, PartialEq)]
110pub struct DocumentInput {
111    /// Caller-supplied filename used for format detection and returned safely.
112    pub file_name: String,
113    /// Caller-supplied media type used as a format-detection fallback.
114    pub content_type: String,
115    /// Complete document bytes.
116    pub data: Vec<u8>,
117}
118
119impl fmt::Debug for DocumentInput {
120    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
121        formatter
122            .debug_struct("DocumentInput")
123            .field("file_name", &self.file_name)
124            .field("content_type", &self.content_type)
125            .field("bytes", &self.data.len())
126            .finish()
127    }
128}
129
130/// Successfully extracted normalized text and metadata.
131#[derive(Clone, Debug, Eq, PartialEq)]
132pub struct ExtractedDocument {
133    /// Safe display filename.
134    pub file_name: String,
135    /// Original normalized media type.
136    pub content_type: String,
137    /// Detected document format.
138    pub format: DocumentFormat,
139    /// Normalized, non-empty extracted text.
140    pub text: String,
141    /// Returned Unicode character count.
142    pub characters: usize,
143    /// Whether the character limit truncated the text.
144    pub truncated: bool,
145}
146
147/// Stateless local document extractor with fixed limits.
148#[derive(Clone, Debug)]
149pub struct DocumentExtractor {
150    limits: ExtractionLimits,
151}
152
153impl DocumentExtractor {
154    /// Constructs an extractor after validating nonzero limits.
155    pub fn new(limits: ExtractionLimits) -> Result<Self> {
156        if limits.max_bytes == 0 || limits.max_characters == 0 {
157            return Err(Error::new(
158                ErrorKind::InvalidInput,
159                "document byte and character limits must be greater than zero",
160            ));
161        }
162        Ok(Self { limits })
163    }
164
165    /// Extracts and normalizes one complete document in the current thread.
166    pub fn extract(&self, input: DocumentInput) -> Result<ExtractedDocument> {
167        if input.data.is_empty() || input.data.len() > self.limits.max_bytes {
168            return Err(Error::new(
169                ErrorKind::InvalidInput,
170                format!(
171                    "document must contain between 1 and {} bytes",
172                    self.limits.max_bytes
173                ),
174            ));
175        }
176        let content_type = input
177            .content_type
178            .split(';')
179            .next()
180            .unwrap_or(&input.content_type)
181            .trim()
182            .to_ascii_lowercase();
183        let format = detect_format(&input.file_name, &content_type).ok_or_else(|| {
184            Error::new(
185                ErrorKind::UnsupportedFormat,
186                "supported documents are PDF, DOCX, XLSX, XLS, XLSB, ODS, CSV, TSV, and plain text",
187            )
188        })?;
189        let file_name = safe_filename(&input.file_name, format);
190        let extracted = extract_text(format, &input.data).map_err(|message| {
191            Error::new(
192                ErrorKind::ExtractionFailed,
193                format!("document could not be converted to text: {message}"),
194            )
195        })?;
196        if extracted.is_empty() {
197            let message = if format == DocumentFormat::Pdf {
198                "PDF contains no extractable text and may require OCR"
199            } else {
200                "document contains no extractable text"
201            };
202            return Err(Error::new(ErrorKind::EmptyText, message));
203        }
204        let (text, truncated) = truncate_characters(&extracted, self.limits.max_characters);
205        let characters = text.chars().count();
206        Ok(ExtractedDocument {
207            file_name,
208            content_type,
209            format,
210            text,
211            characters,
212            truncated,
213        })
214    }
215}
216
217impl Default for DocumentExtractor {
218    fn default() -> Self {
219        Self::new(ExtractionLimits::default()).expect("default document limits are valid")
220    }
221}
222
223fn detect_format(file_name: &str, content_type: &str) -> Option<DocumentFormat> {
224    let extension = file_name
225        .rsplit_once('.')
226        .map(|(_, extension)| extension.to_ascii_lowercase());
227    match extension.as_deref() {
228        Some("pdf") => Some(DocumentFormat::Pdf),
229        Some("docx") => Some(DocumentFormat::Docx),
230        Some("xlsx" | "xls" | "xlsb" | "ods") => Some(DocumentFormat::Spreadsheet),
231        Some("csv" | "tsv" | "txt" | "md" | "json" | "yaml" | "yml" | "xml") => {
232            Some(DocumentFormat::PlainText)
233        }
234        _ if content_type == "application/pdf" => Some(DocumentFormat::Pdf),
235        _ if content_type
236            == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" =>
237        {
238            Some(DocumentFormat::Docx)
239        }
240        _ if matches!(
241            content_type,
242            "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
243                | "application/vnd.ms-excel"
244                | "application/vnd.ms-excel.sheet.binary.macroenabled.12"
245                | "application/vnd.oasis.opendocument.spreadsheet"
246        ) =>
247        {
248            Some(DocumentFormat::Spreadsheet)
249        }
250        _ if content_type.starts_with("text/")
251            || matches!(
252                content_type,
253                "application/json" | "application/xml" | "application/yaml" | "application/x-yaml"
254            ) =>
255        {
256            Some(DocumentFormat::PlainText)
257        }
258        _ => None,
259    }
260}
261
262fn safe_filename(value: &str, format: DocumentFormat) -> String {
263    let cleaned = value
264        .chars()
265        .filter(|character| {
266            !character.is_control() && !matches!(character, '/' | '\\' | ':' | '\0')
267        })
268        .take(200)
269        .collect::<String>();
270    if cleaned.trim().is_empty() {
271        format!("document.{}", format.as_str())
272    } else {
273        cleaned
274    }
275}
276
277fn local_xml_name(name: &[u8]) -> &[u8] {
278    name.rsplit(|byte| *byte == b':').next().unwrap_or(name)
279}
280
281fn extract_docx_text(bytes: &[u8]) -> std::result::Result<String, String> {
282    let mut archive =
283        zip::ZipArchive::new(Cursor::new(bytes)).map_err(|error| error.to_string())?;
284    let mut document = archive
285        .by_name("word/document.xml")
286        .map_err(|error| format!("word/document.xml: {error}"))?;
287    let mut xml = String::new();
288    document
289        .read_to_string(&mut xml)
290        .map_err(|error| error.to_string())?;
291    let mut reader = XmlReader::from_str(&xml);
292    let mut output = String::new();
293    let mut in_text = false;
294    let mut in_cell = false;
295    loop {
296        match reader.read_event() {
297            Ok(XmlEvent::Start(event)) => match local_xml_name(event.name().as_ref()) {
298                b"t" => in_text = true,
299                b"tc" => in_cell = true,
300                _ => {}
301            },
302            Ok(XmlEvent::Empty(event)) => match local_xml_name(event.name().as_ref()) {
303                b"tab" => output.push('\t'),
304                b"br" | b"cr" => output.push('\n'),
305                _ => {}
306            },
307            Ok(XmlEvent::Text(text)) if in_text => {
308                let decoded = text.decode().map_err(|error| error.to_string())?;
309                output.push_str(&decoded);
310            }
311            Ok(XmlEvent::CData(text)) if in_text => {
312                let decoded = text.decode().map_err(|error| error.to_string())?;
313                output.push_str(&decoded);
314            }
315            Ok(XmlEvent::GeneralRef(reference)) if in_text => {
316                let decoded = reference.decode().map_err(|error| error.to_string())?;
317                let escaped = format!("&{decoded};");
318                let resolved =
319                    quick_xml::escape::unescape(&escaped).map_err(|error| error.to_string())?;
320                output.push_str(&resolved);
321            }
322            Ok(XmlEvent::End(event)) => match local_xml_name(event.name().as_ref()) {
323                b"t" => in_text = false,
324                b"p" if in_cell => output.push_str(" / "),
325                b"p" => output.push('\n'),
326                b"tc" => {
327                    if output.ends_with(" / ") {
328                        output.truncate(output.len() - 3);
329                    }
330                    output.push('\t');
331                    in_cell = false;
332                }
333                b"tr" => {
334                    if output.ends_with('\t') {
335                        output.pop();
336                    }
337                    output.push('\n');
338                }
339                _ => {}
340            },
341            Ok(XmlEvent::Eof) => break,
342            Err(error) => return Err(error.to_string()),
343            _ => {}
344        }
345    }
346    Ok(output)
347}
348
349fn extract_spreadsheet_text(bytes: &[u8]) -> std::result::Result<String, String> {
350    let mut workbook = open_workbook_auto_from_rs(Cursor::new(bytes.to_vec()))
351        .map_err(|error| error.to_string())?;
352    let sheet_names = workbook.sheet_names().to_vec();
353    let mut output = String::new();
354    for (sheet_index, sheet_name) in sheet_names.iter().enumerate() {
355        if sheet_index > 0 {
356            output.push('\n');
357        }
358        output.push_str("Sheet: ");
359        output.push_str(sheet_name);
360        output.push('\n');
361        let range = workbook
362            .worksheet_range(sheet_name)
363            .map_err(|error| format!("{sheet_name}: {error}"))?;
364        for row in range.rows() {
365            let line = row
366                .iter()
367                .map(ToString::to_string)
368                .map(|cell| cell.replace(['\t', '\r', '\n'], " "))
369                .collect::<Vec<_>>()
370                .join("\t");
371            output.push_str(line.trim_end());
372            output.push('\n');
373        }
374    }
375    Ok(output)
376}
377
378fn normalize_text(value: String) -> String {
379    let normalized = value
380        .replace("\r\n", "\n")
381        .replace('\r', "\n")
382        .replace('\0', "");
383    let mut output = String::with_capacity(normalized.len());
384    let mut blank_lines = 0;
385    for line in normalized.lines() {
386        let line = line.trim_end();
387        if line.trim().is_empty() {
388            blank_lines += 1;
389            if blank_lines > 1 {
390                continue;
391            }
392        } else {
393            blank_lines = 0;
394        }
395        output.push_str(line);
396        output.push('\n');
397    }
398    output.trim().to_owned()
399}
400
401fn extract_text(format: DocumentFormat, bytes: &[u8]) -> std::result::Result<String, String> {
402    let extracted = match format {
403        DocumentFormat::Pdf => {
404            pdf_extract::extract_text_from_mem(bytes).map_err(|error| error.to_string())?
405        }
406        DocumentFormat::Docx => extract_docx_text(bytes)?,
407        DocumentFormat::Spreadsheet => extract_spreadsheet_text(bytes)?,
408        DocumentFormat::PlainText => String::from_utf8_lossy(bytes)
409            .trim_start_matches('\u{feff}')
410            .to_owned(),
411    };
412    Ok(normalize_text(extracted))
413}
414
415fn truncate_characters(value: &str, limit: usize) -> (String, bool) {
416    let mut iter = value.char_indices();
417    let Some((boundary, _)) = iter.nth(limit) else {
418        return (value.to_owned(), false);
419    };
420    (value[..boundary].trim_end().to_owned(), true)
421}
422
423#[cfg(test)]
424mod tests {
425    use super::*;
426
427    #[test]
428    fn formats_are_selected_from_extensions_and_media_types() {
429        assert_eq!(
430            detect_format("report.PDF", "application/octet-stream"),
431            Some(DocumentFormat::Pdf)
432        );
433        assert_eq!(
434            detect_format("book.xlsx", "application/octet-stream"),
435            Some(DocumentFormat::Spreadsheet)
436        );
437        assert_eq!(
438            detect_format("notes", "text/plain; charset=utf-8"),
439            Some(DocumentFormat::PlainText)
440        );
441        assert_eq!(detect_format("archive.zip", "application/zip"), None);
442    }
443
444    #[test]
445    fn plain_text_is_normalized_and_bounded() {
446        let extractor = DocumentExtractor::new(ExtractionLimits {
447            max_bytes: 100,
448            max_characters: 5,
449        })
450        .unwrap();
451        let extracted = extractor
452            .extract(DocumentInput {
453                file_name: "notes.txt".into(),
454                content_type: "text/plain".into(),
455                data: b"hello\r\n\r\n\r\nworld\0".to_vec(),
456            })
457            .unwrap();
458        assert_eq!(extracted.text, "hello");
459        assert!(extracted.truncated);
460    }
461}
kcode_doc_extraction/lib.rs

kcode_doc_extraction/
lib.rs
