use crate::ExtractError;
pub fn extract_tables_impl(
pdf_bytes: &[u8],
page_num: Option<usize>,
) -> Result<Vec<Vec<Vec<String>>>, ExtractError> {
let page_texts = extract_page_texts(pdf_bytes)?;
let pages_to_process: Vec<&str> = match page_num {
Some(n) if n < page_texts.len() => vec![page_texts[n].as_str()],
Some(_) => return Ok(vec![]),
None => page_texts.iter().map(String::as_str).collect(),
};
let mut all_tables: Vec<Vec<Vec<String>>> = Vec::new();
for page_text in pages_to_process {
all_tables.extend(detect_tables_in_page(page_text));
}
Ok(all_tables)
}
fn extract_page_texts(pdf_bytes: &[u8]) -> Result<Vec<String>, ExtractError> {
let doc = crate::document::open_sp_with_password(pdf_bytes, b"")?;
let mut page_nums: Vec<u32> = doc.get_pages().into_keys().collect();
page_nums.sort();
Ok(page_nums
.iter()
.map(|&n| doc.extract_text(&[n]).unwrap_or_default())
.collect())
}
fn detect_tables_in_page(page_text: &str) -> Vec<Vec<Vec<String>>> {
let lines: Vec<&str> = page_text.lines().collect();
if lines.is_empty() {
return vec![];
}
let regions = identify_table_regions(&lines);
let mut tables = Vec::new();
for region in regions {
if region.len() < 2 {
continue;
}
let table = parse_table_region(®ion);
if !table.is_empty() && table.len() >= 2 {
let has_multi_col = table.iter().any(|row| row.len() >= 2);
if has_multi_col {
tables.push(table);
}
}
}
tables
}
fn identify_table_regions(lines: &[&str]) -> Vec<Vec<String>> {
let mut regions: Vec<Vec<String>> = Vec::new();
let mut current_region: Vec<String> = Vec::new();
let mut consecutive_structured = 0usize;
for &line in lines {
let stripped = line.trim();
if stripped.is_empty() {
if consecutive_structured >= 2 && !current_region.is_empty() {
regions.push(current_region.clone());
current_region.clear();
}
consecutive_structured = 0;
continue;
}
let has_gaps = has_column_gaps(stripped);
if has_gaps {
consecutive_structured += 1;
current_region.push(line.to_string());
} else {
if consecutive_structured >= 1 {
current_region.push(line.to_string());
} else {
if !current_region.is_empty() {
if consecutive_structured >= 2 {
regions.push(current_region.clone());
}
current_region.clear();
}
consecutive_structured = 0;
}
}
}
if consecutive_structured >= 2 && !current_region.is_empty() {
regions.push(current_region);
}
regions
}
fn has_column_gaps(line: &str) -> bool {
let mut gap_count = 0usize;
let chars: Vec<char> = line.chars().collect();
let n = chars.len();
let mut i = 0;
while i < n {
if chars[i] == ' ' {
let gap_start = i;
while i < n && chars[i] == ' ' {
i += 1;
}
let gap_len = i - gap_start;
if gap_len >= 2 {
gap_count += 1;
}
} else {
i += 1;
}
}
gap_count >= 1
}
fn parse_table_region(lines: &[String]) -> Vec<Vec<String>> {
if lines.is_empty() {
return vec![];
}
let col_breaks = find_column_breaks(lines);
lines
.iter()
.map(|line| split_line_at_breaks(line, &col_breaks))
.filter(|row| !row.is_empty() && row.iter().any(|c| !c.is_empty()))
.collect()
}
fn find_column_breaks(lines: &[String]) -> Vec<usize> {
let lines_chars: Vec<Vec<char>> = lines.iter().map(|l| l.chars().collect()).collect();
let max_chars = lines_chars.iter().map(|c| c.len()).max().unwrap_or(0);
if max_chars == 0 {
return vec![];
}
let mut space_votes: Vec<usize> = vec![0; max_chars + 1];
let n_lines = lines_chars.len();
for chars in &lines_chars {
for (i, &ch) in chars.iter().enumerate() {
if ch == ' ' {
space_votes[i] += 1;
}
}
for vote in &mut space_votes[chars.len()..max_chars] {
*vote += 1;
}
}
let threshold = (n_lines as f64 * 0.6) as usize;
let gap_mask: Vec<bool> = space_votes.iter().map(|&v| v >= threshold).collect();
let mut breaks: Vec<usize> = Vec::new();
let mut in_gap = false;
let mut gap_start = 0;
for (i, &is_gap) in gap_mask.iter().enumerate() {
if is_gap && !in_gap {
in_gap = true;
gap_start = i;
} else if !is_gap && in_gap {
in_gap = false;
let gap_len = i - gap_start;
if gap_len >= 2 {
breaks.push(i);
}
}
}
if in_gap {
let gap_len = max_chars - gap_start;
if gap_len >= 2 {
breaks.push(max_chars);
}
}
breaks
}
fn split_line_at_breaks(line: &str, breaks: &[usize]) -> Vec<String> {
if breaks.is_empty() {
let t = line.trim().to_string();
return if t.is_empty() { vec![] } else { vec![t] };
}
let chars: Vec<char> = line.chars().collect();
let mut cells: Vec<String> = Vec::new();
let mut prev = 0;
for &brk in breaks {
let end = brk.min(chars.len());
let cell: String = chars[prev..end].iter().collect();
cells.push(cell.trim().to_string());
prev = end;
}
if prev < chars.len() {
let cell: String = chars[prev..].iter().collect();
cells.push(cell.trim().to_string());
}
cells
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_column_gap_detection() {
let line = "Moody's Aa2";
assert!(has_column_gaps(line));
}
#[test]
fn test_split_at_breaks() {
let line = "Moody's Aa2".to_string();
let breaks = find_column_breaks(&[
"Agency Rating".to_string(),
"Moody's Aa2".to_string(),
"S&P AA+".to_string(),
]);
assert!(!breaks.is_empty(), "Should detect column breaks");
let cells = split_line_at_breaks(&line, &breaks);
assert!(cells.len() >= 2, "Should have at least 2 cells");
}
#[test]
fn test_simple_table_region() {
let lines = vec![
"Agency Rating Outlook".to_string(),
"Moody's Aa2 Stable".to_string(),
"S&P AA+ Stable".to_string(),
"Fitch AA Stable".to_string(),
];
let table = parse_table_region(&lines);
assert!(!table.is_empty(), "Should parse table");
assert!(table[0].len() >= 2, "Should have multiple columns");
}
}