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, /// All Calendar by `service_id` pub calendar: HashMap, /// All calendar dates grouped by service_id pub calendar_dates: HashMap>, /// All routes by `route_id` pub routes: HashMap, /// All stop by `stop_id`. pub stops: HashMap, } // Utility function to load all records in a dataset fn load_all(_: &V) -> bool { true } // Loads a vector of the selected type fn load_vector( destination: &mut Vec, zip_reader: &mut ZipArchive, 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, zip_reader: &mut ZipArchive, 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>, zip_reader: &mut ZipArchive, 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) -> 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, 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; }