115 lines
4.2 KiB
Rust
115 lines
4.2 KiB
Rust
mod loader;
|
|
mod utils;
|
|
pub mod structs;
|
|
use chrono::{DateTime, Local};
|
|
use std::{
|
|
collections::{HashMap, HashSet},
|
|
fs::File,
|
|
};
|
|
use gtfs_structures::{Agency, Calendar, CalendarDate, Exception, RawStopTime, RawTrip, Route, Stop};
|
|
|
|
use crate::gtfs::{loader::load_gtfs, structs::{Arrival, Preferences}};
|
|
|
|
// 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>,
|
|
/// All trips by trip_id
|
|
pub trips: HashMap<String, RawTrip>,
|
|
/// Stop times for the chosen stops and the chosen routes
|
|
pub stop_times: HashMap<(String, u32), RawStopTime>,
|
|
}
|
|
|
|
|
|
impl Gtfs {
|
|
|
|
pub fn get_next_arrivals_for(&self, target_datetime: &DateTime<Local>) -> Box<Vec<Arrival<'_>>> {
|
|
let naive_target = target_datetime.naive_local();
|
|
let target_date = naive_target.date();
|
|
|
|
// Find which calendars apply
|
|
let mut active_service_ids: HashSet<String> = HashSet::new();
|
|
for (id, calendar) in self.calendar.iter() {
|
|
if calendar.valid_weekday(target_date)
|
|
&& calendar.start_date <= target_date
|
|
&& calendar.end_date > target_date {
|
|
active_service_ids.insert(id.to_string());
|
|
}
|
|
}
|
|
|
|
// Are there any exceptions for the calendars above?
|
|
for (_calendar_id, exceptions) in self.calendar_dates.iter() {
|
|
for exception in exceptions.iter() {
|
|
if exception.date.eq(&target_date) {
|
|
match exception.exception_type {
|
|
Exception::Added => {
|
|
active_service_ids.insert(exception.service_id.clone());
|
|
}
|
|
Exception::Deleted => {
|
|
active_service_ids.remove(&exception.service_id);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Find the trips happening on these calendars
|
|
let mut trips: HashMap<&String, &RawTrip> = HashMap::new();
|
|
for (_id, trip) in self.trips.iter() {
|
|
if active_service_ids.contains(&trip.service_id) {
|
|
trips.insert(&trip.id, trip);
|
|
}
|
|
}
|
|
|
|
// Finally, find the arrivals for the active trips on the chosen stops
|
|
let mut arrivals: Vec<Arrival> = Vec::new();
|
|
let current_timestamp = target_datetime.timestamp();
|
|
for (_id, stop_time) in self.stop_times.iter() {
|
|
if trips.contains_key(&stop_time.trip_id) {
|
|
let stop_timestamp = stop_time.departure_time.or(stop_time.arrival_time).unwrap();
|
|
if current_timestamp < stop_timestamp.into() {
|
|
let arrival: Arrival = Arrival {
|
|
route: self.routes.get(&self.trips.get(&stop_time.trip_id).unwrap().route_id).unwrap(),
|
|
stop: self.stops.get(&stop_time.stop_id).unwrap(),
|
|
departure_time: stop_timestamp
|
|
};
|
|
arrivals.push(arrival);
|
|
}
|
|
}
|
|
}
|
|
|
|
return Box::from(arrivals);
|
|
}
|
|
|
|
|
|
/// Load a GTFS structure from a zip file
|
|
pub fn load(src_file: &str, prefs: &Preferences) -> 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(),
|
|
trips: HashMap::new(),
|
|
stop_times: HashMap::new(),
|
|
};
|
|
|
|
|
|
load_gtfs(&mut gtfs, &mut zip_reader, &prefs.route_numbers, &prefs.stop_codes);
|
|
|
|
return gtfs;
|
|
}
|
|
}
|