133 lines
3.9 KiB
Rust
133 lines
3.9 KiB
Rust
use gtfs_structures::{Agency, Calendar, CalendarDate, Route, Stop};
|
|
use serde::{de::DeserializeOwned};
|
|
use std::{collections::HashMap, fs::File};
|
|
use zip::ZipArchive;
|
|
|
|
// The main GTFS struct. This is similar to (but not exactly) gtfs-structures::Gtfs because we don't need everything
|
|
#[derive(Debug)]
|
|
pub struct Gtfs {
|
|
/// All agencies. They can not be read by `agency_id`, as it is not a required field
|
|
pub agencies: Vec<Agency>,
|
|
/// All Calendar by `service_id`
|
|
pub calendar: HashMap<String, Calendar>,
|
|
/// All calendar dates grouped by service_id
|
|
pub calendar_dates: HashMap<String, Vec<CalendarDate>>,
|
|
/// All routes by `route_id`
|
|
pub routes: HashMap<String, Route>,
|
|
/// All stop by `stop_id`.
|
|
pub stops: HashMap<String, Stop>,
|
|
}
|
|
|
|
// Utility function to load all records in a dataset
|
|
fn load_all<V>(_: &V) -> bool { true }
|
|
|
|
// Loads a vector of the selected type
|
|
fn load_vector<T: serde::de::DeserializeOwned>(
|
|
destination: &mut Vec<T>,
|
|
zip_reader: &mut ZipArchive<File>,
|
|
table_name: &str,
|
|
) {
|
|
let file_reader = zip_reader.by_name(table_name).unwrap();
|
|
let mut rdr = csv::Reader::from_reader(file_reader);
|
|
|
|
for row in rdr.deserialize() {
|
|
let record: T = row.unwrap();
|
|
destination.push(record);
|
|
}
|
|
}
|
|
|
|
// Loads a HashMap of the selected type, using the provided index function as the key
|
|
fn load_map<'a, V: DeserializeOwned>(
|
|
destination: &mut HashMap<String, V>,
|
|
zip_reader: &mut ZipArchive<File>,
|
|
table_name: &str,
|
|
index: fn(&V) -> String,
|
|
accept: fn(&V) -> bool,
|
|
) {
|
|
let file_reader = zip_reader.by_name(table_name).unwrap();
|
|
let mut rdr = csv::Reader::from_reader(file_reader);
|
|
|
|
for row in rdr.deserialize() {
|
|
let record: V = row.unwrap();
|
|
if accept(&record) {
|
|
let idx: String = index(&record);
|
|
destination.insert(idx, record);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Loads a HashMap of a vector of the selected type, using the provided index function as the key
|
|
// And a predicate as a filter
|
|
fn load_vector_map<'a, V: DeserializeOwned + Clone>(
|
|
destination: &mut HashMap<String, Vec<V>>,
|
|
zip_reader: &mut ZipArchive<File>,
|
|
table_name: &str,
|
|
index: fn(&V) -> String,
|
|
accept: fn(&V) -> bool,
|
|
) {
|
|
let file_reader = zip_reader.by_name(table_name).unwrap();
|
|
let mut rdr = csv::Reader::from_reader(file_reader);
|
|
|
|
for row in rdr.deserialize() {
|
|
let record: V = row.unwrap();
|
|
if accept(&record) {
|
|
let idx: String = index(&record);
|
|
destination.entry(idx).or_insert_with(Vec::new).push(record);
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn init(src_file: &str, routes: Vec<String>) -> Gtfs {
|
|
// Open zip file
|
|
let mut zip_reader = zip::ZipArchive::new(File::open(src_file).unwrap()).unwrap();
|
|
|
|
let mut gtfs: Gtfs = Gtfs {
|
|
agencies: Vec::new(),
|
|
calendar: HashMap::new(),
|
|
calendar_dates: HashMap::new(),
|
|
routes: HashMap::new(),
|
|
stops: HashMap::new(),
|
|
};
|
|
|
|
// Agencies
|
|
load_vector(&mut gtfs.agencies, &mut zip_reader, "agency.txt");
|
|
|
|
// Calendars
|
|
load_map(
|
|
&mut gtfs.calendar,
|
|
&mut zip_reader,
|
|
"calendar.txt",
|
|
|c: &Calendar| String::from(&c.id),
|
|
load_all
|
|
);
|
|
|
|
// Calendar Dates
|
|
load_vector_map(
|
|
&mut gtfs.calendar_dates,
|
|
&mut zip_reader,
|
|
"calendar_dates.txt",
|
|
|d: &CalendarDate| String::from(&d.service_id),
|
|
load_all
|
|
);
|
|
|
|
// Stops
|
|
load_map(&mut gtfs.stops,
|
|
&mut zip_reader,
|
|
"stops.txt",
|
|
|s: &Stop| {String::from(&s.id)},
|
|
load_all
|
|
);
|
|
|
|
// Routes
|
|
let accept_filter: fn(&Route) -> bool = (| rs: Vec<String>, r: &Route | {rs.contains(&r.short_name.unwrap())}).curry(routes);
|
|
load_map(
|
|
&mut gtfs.routes,
|
|
&mut zip_reader,
|
|
"routes.txt",
|
|
|r: &Route| String::from(&r.id),
|
|
accept_filter
|
|
);
|
|
|
|
return gtfs;
|
|
}
|