factor out extraction code into a module
This commit is contained in:
parent
7b22bbab16
commit
e640a7cd09
212
src/extract.rs
Normal file
212
src/extract.rs
Normal file
|
|
@ -0,0 +1,212 @@
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::fmt::Write;
|
||||||
|
use std::rc::Rc;
|
||||||
|
|
||||||
|
use lazy_static::lazy_static;
|
||||||
|
use pdf::{backend::Backend, content::Operation, primitive::Primitive};
|
||||||
|
use regex::Regex;
|
||||||
|
use time::Date;
|
||||||
|
|
||||||
|
const POSITION_ERROR_MARGIN: f32 = 2.0;
|
||||||
|
|
||||||
|
lazy_static! {
|
||||||
|
static ref DATE_REGEX: Regex = Regex::new(r"^\d{1,2}/\d{1,2}/\d{4}$").unwrap();
|
||||||
|
static ref VALUE_REGEX: Regex = Regex::new(r"^\d+$").unwrap();
|
||||||
|
static ref WHITESPACE_REGEX: Regex = Regex::new(r"\s+").unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
const DATE_PARSE_FORMAT: &[time::format_description::FormatItem] =
|
||||||
|
time::macros::format_description!("[month padding:none]/[day padding:none]/[year]");
|
||||||
|
const DATE_DISPLAY_FORMAT: &[time::format_description::FormatItem] =
|
||||||
|
time::macros::format_description!("[year]-[month]-[day]");
|
||||||
|
|
||||||
|
pub struct DataSet {
|
||||||
|
pub columns: Vec<Rc<String>>,
|
||||||
|
pub rows: Vec<DataPoint>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl DataSet {
|
||||||
|
pub fn extract<B: Backend>(doc: &pdf::file::File<B>) -> Option<Self> {
|
||||||
|
let mut doc_iter = DocumentIterator::new(doc).peekable();
|
||||||
|
|
||||||
|
let mut columns: Vec<(Rc<String>, f32)> = Vec::new();
|
||||||
|
let mut rows: Vec<DataPoint> = Vec::new();
|
||||||
|
|
||||||
|
let (mut current_datapoint, mut current_y) = loop {
|
||||||
|
if let Some(text) = doc_iter.next() {
|
||||||
|
if is_new_column_header(&text.text) {
|
||||||
|
let mut column_name = text.text;
|
||||||
|
let column_x = text.x;
|
||||||
|
while let Some(more) = doc_iter.peek() {
|
||||||
|
if is_new_column_header(&more.text) || DATE_REGEX.is_match(&more.text) {
|
||||||
|
columns.push((Rc::new(column_name), column_x));
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
column_name += " ";
|
||||||
|
column_name += &more.text;
|
||||||
|
doc_iter.next();
|
||||||
|
}
|
||||||
|
} else if DATE_REGEX.is_match(&text.text) {
|
||||||
|
break (
|
||||||
|
DataPoint::new(&text).expect("Failed to parse date!"),
|
||||||
|
text.y,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
columns.sort_by(|(_, x0), (_, x1)| x0.partial_cmp(x1).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
|
|
||||||
|
for text in doc_iter {
|
||||||
|
if DATE_REGEX.is_match(&text.text) {
|
||||||
|
rows.push(std::mem::replace(
|
||||||
|
&mut current_datapoint,
|
||||||
|
DataPoint::new(&text).expect("Failed to parse date!"),
|
||||||
|
));
|
||||||
|
current_y = text.y;
|
||||||
|
} else if VALUE_REGEX.is_match(&text.text) {
|
||||||
|
if (current_y - text.y).abs() > POSITION_ERROR_MARGIN {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if let Some((column, _)) = columns.iter().rev().find(|(_, x)| *x < text.x) {
|
||||||
|
current_datapoint.values.insert(
|
||||||
|
column.clone(),
|
||||||
|
text.text.parse().expect("Failed to parse value!"),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Some(Self {
|
||||||
|
columns: columns.into_iter().map(|(column, _)| column).collect(),
|
||||||
|
rows,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn csv_rows(&self) -> impl Iterator<Item = Result<String, std::fmt::Error>> + '_ {
|
||||||
|
std::iter::once_with(|| {
|
||||||
|
let mut header = String::from("Date");
|
||||||
|
for column in self.columns.iter() {
|
||||||
|
write!(&mut header, ",{}", column)?;
|
||||||
|
}
|
||||||
|
Ok(header)
|
||||||
|
})
|
||||||
|
.chain(self.rows.iter().map(|datapoint| {
|
||||||
|
let mut csv_row = datapoint
|
||||||
|
.date
|
||||||
|
.format(DATE_DISPLAY_FORMAT)
|
||||||
|
.expect("Failed to format date!");
|
||||||
|
for column in self.columns.iter() {
|
||||||
|
if let Some(val) = datapoint.values.get(column) {
|
||||||
|
write!(&mut csv_row, ",{}", val)?;
|
||||||
|
} else {
|
||||||
|
write!(&mut csv_row, ",")?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(csv_row)
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct DataPoint {
|
||||||
|
pub date: Date,
|
||||||
|
pub values: HashMap<Rc<String>, u32>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl DataPoint {
|
||||||
|
fn new(text: &DocumentText) -> Result<Self, time::error::Parse> {
|
||||||
|
Ok(Self {
|
||||||
|
date: Date::parse(&text.text, DATE_PARSE_FORMAT)?,
|
||||||
|
values: HashMap::new(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct DocumentText {
|
||||||
|
x: f32,
|
||||||
|
y: f32,
|
||||||
|
text: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct DocumentIterator<'a> {
|
||||||
|
x: f32,
|
||||||
|
y: f32,
|
||||||
|
operations: Box<dyn Iterator<Item = Operation> + 'a>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'a> DocumentIterator<'a> {
|
||||||
|
fn new<B: Backend>(document: &'a pdf::file::File<B>) -> Self {
|
||||||
|
Self {
|
||||||
|
x: 0.0,
|
||||||
|
y: 0.0,
|
||||||
|
operations: Box::new(
|
||||||
|
document
|
||||||
|
.pages()
|
||||||
|
.filter_map(|page| Some(page.ok()?.contents.clone()?.operations.into_iter()))
|
||||||
|
.flatten(),
|
||||||
|
),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'a> Iterator for DocumentIterator<'a> {
|
||||||
|
type Item = DocumentText;
|
||||||
|
|
||||||
|
fn next(&mut self) -> Option<Self::Item> {
|
||||||
|
for Operation { operator, operands } in self.operations.as_mut() {
|
||||||
|
if operator == "Tm" {
|
||||||
|
if let (Some(x), Some(y)) = (
|
||||||
|
operands.get(4).and_then(extract_number),
|
||||||
|
operands.get(5).and_then(extract_number),
|
||||||
|
) {
|
||||||
|
self.x = x;
|
||||||
|
self.y = y;
|
||||||
|
}
|
||||||
|
} else if operator == "TJ" || operator == "Tj" {
|
||||||
|
if let Some(text) = operands.get(0).and_then(extract_string) {
|
||||||
|
return Some(DocumentText {
|
||||||
|
x: self.x,
|
||||||
|
y: self.y,
|
||||||
|
text,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn extract_number(p: &Primitive) -> Option<f32> {
|
||||||
|
match p {
|
||||||
|
Primitive::Number(n) => Some(*n),
|
||||||
|
Primitive::Integer(n) => Some(*n as f32),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn extract_string(p: &Primitive) -> Option<String> {
|
||||||
|
let result: Box<dyn AsRef<str>> = match p {
|
||||||
|
Primitive::Array(array) => {
|
||||||
|
let mut acc = String::new();
|
||||||
|
for element in array.iter() {
|
||||||
|
if let Primitive::String(s) = element {
|
||||||
|
acc += s.as_str().ok()?.as_ref();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Box::new(acc)
|
||||||
|
}
|
||||||
|
Primitive::String(s) => Box::new(s.as_str().ok()?),
|
||||||
|
_ => return None,
|
||||||
|
};
|
||||||
|
Some(
|
||||||
|
WHITESPACE_REGEX
|
||||||
|
.replace_all((*result).as_ref().trim(), " ")
|
||||||
|
.to_string(),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn is_new_column_header(s: &str) -> bool {
|
||||||
|
s.starts_with("Northern") || s.starts_with("Southern")
|
||||||
|
}
|
||||||
189
src/main.rs
189
src/main.rs
|
|
@ -1,190 +1,11 @@
|
||||||
use std::collections::HashMap;
|
mod extract;
|
||||||
|
|
||||||
use lazy_static::lazy_static;
|
use extract::DataSet;
|
||||||
use pdf::{primitive::Primitive, content::Operation, backend::Backend};
|
|
||||||
use regex::Regex;
|
|
||||||
use time::Date;
|
|
||||||
|
|
||||||
const POSITION_ERROR_MARGIN: f32 = 2.0;
|
|
||||||
|
|
||||||
lazy_static! {
|
|
||||||
static ref DATE_REGEX: Regex = Regex::new(r"^\d{1,2}/\d{1,2}/\d{4}$").unwrap();
|
|
||||||
static ref VALUE_REGEX: Regex = Regex::new(r"^\d+$").unwrap();
|
|
||||||
static ref WHITESPACE_REGEX: Regex = Regex::new(r"\s+").unwrap();
|
|
||||||
}
|
|
||||||
|
|
||||||
const DATE_PARSE_FORMAT: &[time::format_description::FormatItem] = time::macros::format_description!(
|
|
||||||
"[month padding:none]/[day padding:none]/[year]"
|
|
||||||
);
|
|
||||||
const DATE_DISPLAY_FORMAT: &[time::format_description::FormatItem] = time::macros::format_description!(
|
|
||||||
"[year]-[month]-[day]"
|
|
||||||
);
|
|
||||||
|
|
||||||
#[derive(Debug)]
|
|
||||||
struct DataPoint<'a> {
|
|
||||||
date: Date,
|
|
||||||
values: HashMap<&'a str, u32>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl<'a> DataPoint<'a> {
|
|
||||||
fn new(text: &DocumentText) -> Result<Self, time::error::Parse> {
|
|
||||||
Ok(Self { date: Date::parse(&text.text, DATE_PARSE_FORMAT)?, values: HashMap::new() })
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
struct DocumentText {
|
|
||||||
x: f32,
|
|
||||||
y: f32,
|
|
||||||
text: String,
|
|
||||||
}
|
|
||||||
|
|
||||||
struct DocumentIterator<'a> {
|
|
||||||
x: f32,
|
|
||||||
y: f32,
|
|
||||||
operations: Box<dyn Iterator<Item = Operation> + 'a>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl<'a> DocumentIterator<'a> {
|
|
||||||
fn new<B: Backend>(document: &'a pdf::file::File<B>) -> Self {
|
|
||||||
Self {
|
|
||||||
x: 0.0,
|
|
||||||
y: 0.0,
|
|
||||||
operations: Box::new(
|
|
||||||
document
|
|
||||||
.pages()
|
|
||||||
.filter_map(|page| {
|
|
||||||
Some(page.ok()?.contents.clone()?.operations.into_iter())
|
|
||||||
})
|
|
||||||
.flatten()
|
|
||||||
),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl<'a> Iterator for DocumentIterator<'a> {
|
|
||||||
type Item = DocumentText;
|
|
||||||
|
|
||||||
fn next(&mut self) -> Option<Self::Item> {
|
|
||||||
for Operation { operator, operands } in self.operations.as_mut() {
|
|
||||||
if operator == "Tm" {
|
|
||||||
if let (Some(x), Some(y)) = (
|
|
||||||
operands.get(4).and_then(extract_number),
|
|
||||||
operands.get(5).and_then(extract_number),
|
|
||||||
) {
|
|
||||||
self.x = x;
|
|
||||||
self.y = y;
|
|
||||||
}
|
|
||||||
} else if operator == "TJ" || operator == "Tj" {
|
|
||||||
if let Some(text) = operands.get(0).and_then(extract_string) {
|
|
||||||
return Some(DocumentText {
|
|
||||||
x: self.x,
|
|
||||||
y: self.y,
|
|
||||||
text,
|
|
||||||
});
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
None
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn extract_number(p: &Primitive) -> Option<f32> {
|
|
||||||
match p {
|
|
||||||
Primitive::Number(n) => Some(*n),
|
|
||||||
Primitive::Integer(n) => Some(*n as f32),
|
|
||||||
_ => None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn extract_string(p: &Primitive) -> Option<String> {
|
|
||||||
let result: Box<dyn AsRef<str>> = match p {
|
|
||||||
Primitive::Array(array) => {
|
|
||||||
let mut acc = String::new();
|
|
||||||
for element in array.iter() {
|
|
||||||
if let Primitive::String(s) = element {
|
|
||||||
acc += s.as_str().ok()?.as_ref();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Box::new(acc)
|
|
||||||
}
|
|
||||||
Primitive::String(s) => Box::new(s.as_str().ok()?),
|
|
||||||
_ => return None
|
|
||||||
};
|
|
||||||
Some(WHITESPACE_REGEX.replace_all((*result).as_ref().trim(), " ").to_string())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn is_new_column_header(s: &str) -> bool {
|
|
||||||
s.starts_with("Northern") || s.starts_with("Southern")
|
|
||||||
}
|
|
||||||
|
|
||||||
fn main() {
|
fn main() {
|
||||||
let doc = pdf::file::File::open("data.pdf").expect("Failed to read PDF");
|
let doc = pdf::file::File::open("data.pdf").expect("Failed to read PDF");
|
||||||
let mut doc_iter = DocumentIterator::new(&doc).peekable();
|
let dataset = DataSet::extract(&doc).expect("Failed to extract dataset");
|
||||||
|
for row in dataset.csv_rows() {
|
||||||
let mut columns: Vec<(String, f32)> = Vec::new();
|
println!("{}", row.unwrap());
|
||||||
let mut datapoints: Vec<DataPoint> = Vec::new();
|
|
||||||
|
|
||||||
let (mut current_datapoint, mut current_y) = loop {
|
|
||||||
if let Some(text) = doc_iter.next() {
|
|
||||||
if is_new_column_header(&text.text) {
|
|
||||||
let mut column_name = text.text;
|
|
||||||
let column_x = text.x;
|
|
||||||
while let Some(more) = doc_iter.peek() {
|
|
||||||
if is_new_column_header(&more.text) || DATE_REGEX.is_match(&more.text) {
|
|
||||||
columns.push((column_name, column_x));
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
column_name += " ";
|
|
||||||
column_name += &more.text;
|
|
||||||
doc_iter.next();
|
|
||||||
}
|
|
||||||
} else if DATE_REGEX.is_match(&text.text) {
|
|
||||||
break (DataPoint::new(&text).expect("Failed to parse date!"), text.y);
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
columns.sort_by(|(_, x0), (_, x1)| x0.partial_cmp(x1).unwrap_or(std::cmp::Ordering::Equal));
|
|
||||||
|
|
||||||
for text in doc_iter {
|
|
||||||
if DATE_REGEX.is_match(&text.text) {
|
|
||||||
datapoints.push(std::mem::replace(
|
|
||||||
&mut current_datapoint,
|
|
||||||
DataPoint::new(&text).expect("Failed to parse date!"),
|
|
||||||
));
|
|
||||||
current_y = text.y;
|
|
||||||
} else if VALUE_REGEX.is_match(&text.text) {
|
|
||||||
if (current_y - text.y).abs() > POSITION_ERROR_MARGIN {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
if let Some((column, _)) = columns.iter().rev().find(|(_, x)| *x < text.x) {
|
|
||||||
current_datapoint.values.insert(column, text.text.parse().expect("Failed to parse value!"));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
print!("Date");
|
|
||||||
for (column, _) in columns.iter() {
|
|
||||||
print!(",{}", column);
|
|
||||||
}
|
|
||||||
println!();
|
|
||||||
for datapoint in datapoints.iter() {
|
|
||||||
print!(
|
|
||||||
"{}",
|
|
||||||
datapoint
|
|
||||||
.date
|
|
||||||
.format(DATE_DISPLAY_FORMAT)
|
|
||||||
.expect("Failed to format date!")
|
|
||||||
);
|
|
||||||
for (column, _) in columns.iter() {
|
|
||||||
if let Some(val) = datapoint.values.get(&column.as_ref()) {
|
|
||||||
print!(",{}", val);
|
|
||||||
} else {
|
|
||||||
print!(",");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
println!();
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue