gtfs_structures/
gtfs_reader.rs

1use serde::Deserialize;
2use sha2::{Digest, Sha256};
3
4use crate::{Error, Gtfs, RawGtfs};
5use std::collections::HashMap;
6use std::convert::TryFrom;
7use std::fs::File;
8use std::io::Read;
9use std::path::Path;
10use web_time::Instant;
11
12/// Allows to parameterize how the parsing library behaves
13///
14/// ```
15///let gtfs = gtfs_structures::GtfsReader::default()
16///    .read_stop_times(false) // Won’t read the stop times to save time and memory
17///    .read_shapes(false) // Won’t read shapes to save time and memory
18///    .unkown_enum_as_default(false) // Won’t convert unknown enumerations into default (e.g. LocationType=42 considered as a stop point)
19///    .read("fixtures/zips/gtfs.zip")?;
20///assert_eq!(0, gtfs.trips.get("trip1").unwrap().stop_times.len());
21/// # Ok::<(), gtfs_structures::error::Error>(())
22///```
23///
24/// You can also get a [RawGtfs] by doing
25/// ```
26///let gtfs = gtfs_structures::GtfsReader::default()
27///    .read_stop_times(false)
28///    .raw()
29///    .read("fixtures/zips/gtfs.zip")?;
30///assert_eq!(1, gtfs.trips?.len());
31///assert_eq!(0, gtfs.stop_times?.len());
32/// # Ok::<(), gtfs_structures::error::Error>(())
33///```
34pub struct GtfsReader {
35    /// [crate::objects::StopTime] are very large and not always needed. This allows to skip reading them
36    pub read_stop_times: bool,
37    /// [crate::objects::Shape] are very large and not always needed. This allows to skip reading them
38    pub read_shapes: bool,
39    /// If a an enumeration has an unknown value, should we use the default value
40    pub unkown_enum_as_default: bool,
41    /// Avoid trimming the fields
42    ///
43    /// It is quite time consuming
44    /// If performance is an issue, and if your data is high quality, you can switch it off
45    pub trim_fields: bool,
46}
47
48impl Default for GtfsReader {
49    fn default() -> Self {
50        GtfsReader {
51            read_stop_times: true,
52            read_shapes: true,
53            unkown_enum_as_default: false,
54            trim_fields: true,
55        }
56    }
57}
58
59impl GtfsReader {
60    /// Configures the reader to read or not the stop times (default: true)
61    ///
62    /// This can be useful to save time and memory with large datasets when the timetable are not needed
63    /// Returns Self and can be chained
64    pub fn read_stop_times(mut self, read_stop_times: bool) -> Self {
65        self.read_stop_times = read_stop_times;
66        self
67    }
68
69    /// This can be useful to save time and memory with large datasets when shapes are not needed
70    /// Returns Self and can be chained
71    pub fn read_shapes(mut self, read_shapes: bool) -> Self {
72        self.read_shapes = read_shapes;
73        self
74    }
75
76    /// If a an enumeration has un unknown value, should we use the default value (default: false)
77    ///
78    /// For instance, if [crate::objects::Stop] has a [crate::objects::LocationType] with a value 42 in the GTFS
79    /// when true, we will parse it as StopPoint
80    /// when false, we will parse it as Unknown(42)
81    /// Returns Self and can be chained
82    pub fn unkown_enum_as_default(mut self, unkown_enum_as_default: bool) -> Self {
83        self.unkown_enum_as_default = unkown_enum_as_default;
84        self
85    }
86
87    /// Should the fields be trimmed (default: true)
88    ///
89    /// It is quite time consuming
90    /// If performance is an issue, and if your data is high quality, you can set it to false
91    pub fn trim_fields(mut self, trim_fields: bool) -> Self {
92        self.trim_fields = trim_fields;
93        self
94    }
95
96    /// Reads from an url (if starts with `"http"`), or a local path (either a directory or zipped file)
97    ///
98    /// To read from an url, build with read-url feature
99    /// See also [Gtfs::from_url] and [Gtfs::from_path] if you don’t want the library to guess
100    #[cfg(not(target_arch = "wasm32"))]
101    pub fn read(self, gtfs: &str) -> Result<Gtfs, Error> {
102        self.raw().read(gtfs).and_then(Gtfs::try_from)
103    }
104
105    /// Reads the raw GTFS from a local zip archive or local directory
106    pub fn read_from_path<P>(self, path: P) -> Result<Gtfs, Error>
107    where
108        P: AsRef<Path>,
109    {
110        self.raw().read_from_path(path).and_then(Gtfs::try_from)
111    }
112
113    /// Reads the GTFS from a remote url
114    ///
115    /// The library must be built with the read-url feature. Not available on WASM targets.
116    #[cfg(all(feature = "read-url", not(target_arch = "wasm32")))]
117    pub fn read_from_url<U: reqwest::IntoUrl>(self, url: U) -> Result<Gtfs, Error> {
118        self.raw().read_from_url(url).and_then(Gtfs::try_from)
119    }
120
121    /// Asynchronously reads the GTFS from a remote url
122    ///
123    /// The library must be built with the read-url feature
124    #[cfg(feature = "read-url")]
125    pub async fn read_from_url_async<U: reqwest::IntoUrl>(self, url: U) -> Result<Gtfs, Error> {
126        self.raw()
127            .read_from_url_async(url)
128            .await
129            .and_then(Gtfs::try_from)
130    }
131
132    /// Read the Gtfs as a [RawGtfs].
133    ///
134    /// ```
135    ///let gtfs = gtfs_structures::GtfsReader::default()
136    ///    .read_stop_times(false)
137    ///    .raw()
138    ///    .read("fixtures/zips/gtfs.zip")?;
139    ///assert_eq!(1, gtfs.trips?.len());
140    ///assert_eq!(0, gtfs.stop_times?.len());
141    /// # Ok::<(), gtfs_structures::error::Error>(())
142    ///```
143    pub fn raw(self) -> RawGtfsReader {
144        RawGtfsReader { reader: self }
145    }
146}
147
148/// This reader generates [RawGtfs]. It must be built using [GtfsReader::raw]
149///
150/// The methods to read a Gtfs are the same as for [GtfsReader]
151pub struct RawGtfsReader {
152    reader: GtfsReader,
153}
154
155impl RawGtfsReader {
156    fn read_from_directory(&self, p: &std::path::Path) -> Result<RawGtfs, Error> {
157        let start_of_read_instant = Instant::now();
158        // Thoses files are not mandatory
159        // We use None if they don’t exist, not an Error
160        let files = std::fs::read_dir(p)?
161            .filter_map(|d| {
162                d.ok().and_then(|e| {
163                    e.path()
164                        .strip_prefix(p)
165                        .ok()
166                        .and_then(|f| f.to_str().map(|s| s.to_owned()))
167                })
168            })
169            .collect();
170
171        let mut result = RawGtfs {
172            trips: self.read_objs_from_path(p.join("trips.txt")),
173            calendar: self.read_objs_from_optional_path(p, "calendar.txt"),
174            calendar_dates: self.read_objs_from_optional_path(p, "calendar_dates.txt"),
175            stops: self.read_objs_from_path(p.join("stops.txt")),
176            routes: self.read_objs_from_path(p.join("routes.txt")),
177            stop_times: if self.reader.read_stop_times {
178                self.read_objs_from_path(p.join("stop_times.txt"))
179            } else {
180                Ok(Vec::new())
181            },
182            agencies: self.read_objs_from_path(p.join("agency.txt")),
183            shapes: self.read_objs_from_optional_path(p, "shapes.txt"),
184            fare_attributes: self.read_objs_from_optional_path(p, "fare_attributes.txt"),
185            fare_rules: self.read_objs_from_optional_path(p, "fare_rules.txt"),
186            fare_products: self.read_objs_from_optional_path(p, "fare_products.txt"),
187            fare_media: self.read_objs_from_optional_path(p, "fare_media.txt"),
188            rider_categories: self.read_objs_from_optional_path(p, "rider_categories.txt"),
189            frequencies: self.read_objs_from_optional_path(p, "frequencies.txt"),
190            transfers: self.read_objs_from_optional_path(p, "transfers.txt"),
191            pathways: self.read_objs_from_optional_path(p, "pathways.txt"),
192            levels: self.read_objs_from_optional_path(p, "levels.txt"),
193            feed_info: self.read_objs_from_optional_path(p, "feed_info.txt"),
194            read_duration: start_of_read_instant.elapsed(),
195            translations: self.read_objs_from_optional_path(p, "translations.txt"),
196            ticketing_deep_links: self.read_objs_from_optional_path(p, "ticketing_deep_links.txt"),
197            ticketing_identifiers: self
198                .read_objs_from_optional_path(p, "ticketing_identifiers.txt"),
199            attributions: self.read_objs_from_optional_path(p, "attributions.txt"),
200            licenses: self.read_objs_from_optional_path(p, "licenses.txt"),
201            files,
202            source_format: crate::SourceFormat::Directory,
203            sha256: None,
204        };
205
206        if self.reader.unkown_enum_as_default {
207            result.unknown_to_default();
208        }
209        Ok(result)
210    }
211
212    /// Reads from an url (if starts with `"http"`) if the feature `read-url` is activated,
213    /// or a local path (either a directory or zipped file). Not available on WASM targets.
214    #[cfg(not(target_arch = "wasm32"))]
215    pub fn read(self, gtfs: &str) -> Result<RawGtfs, Error> {
216        #[cfg(feature = "read-url")]
217        if gtfs.starts_with("http") {
218            return self.read_from_url(gtfs);
219        }
220        self.read_from_path(gtfs)
221    }
222
223    /// Reads the GTFS from a remote url. Not available on WASM targets.
224    #[cfg(all(feature = "read-url", not(target_arch = "wasm32")))]
225    pub fn read_from_url<U: reqwest::IntoUrl>(self, url: U) -> Result<RawGtfs, Error> {
226        let mut res = reqwest::blocking::get(url)?;
227        let mut body = Vec::new();
228        res.read_to_end(&mut body)?;
229        let cursor = std::io::Cursor::new(body);
230        self.read_from_reader(cursor)
231    }
232
233    /// Asynchronously reads the GTFS from a remote url
234    #[cfg(feature = "read-url")]
235    pub async fn read_from_url_async<U: reqwest::IntoUrl>(self, url: U) -> Result<RawGtfs, Error> {
236        let res = reqwest::get(url).await?.bytes().await?;
237        let reader = std::io::Cursor::new(res);
238        self.read_from_reader(reader)
239    }
240
241    /// Reads the raw GTFS from a local zip archive or local directory
242    pub fn read_from_path<P>(&self, path: P) -> Result<RawGtfs, Error>
243    where
244        P: AsRef<Path>,
245    {
246        let p = path.as_ref();
247        if p.is_file() {
248            let reader = File::open(p)?;
249            self.read_from_reader(reader)
250        } else if p.is_dir() {
251            self.read_from_directory(p)
252        } else {
253            Err(Error::NotFileNorDirectory(format!("{}", p.display())))
254        }
255    }
256
257    pub fn read_from_reader<T: std::io::Read + std::io::Seek>(
258        &self,
259        reader: T,
260    ) -> Result<RawGtfs, Error> {
261        let start_of_read_instant = Instant::now();
262        let hasher = Sha256::new();
263        let mut buf_reader = std::io::BufReader::new(reader);
264        let mut hash_io = digest_io::IoWrapper(hasher);
265        let _n = std::io::copy(&mut buf_reader, &mut hash_io)?;
266        let digest_io::IoWrapper(hasher) = hash_io;
267        let hash = hasher.finalize();
268        let mut archive = zip::ZipArchive::new(buf_reader)?;
269        let mut file_mapping = HashMap::new();
270        let mut files = Vec::new();
271
272        for i in 0..archive.len() {
273            let archive_file = archive.by_index(i)?;
274            files.push(archive_file.name().to_owned());
275
276            for gtfs_file in &[
277                "agency.txt",
278                "calendar.txt",
279                "calendar_dates.txt",
280                "routes.txt",
281                "stops.txt",
282                "stop_times.txt",
283                "trips.txt",
284                "fare_attributes.txt",
285                "fare_rules.txt",
286                "fare_products.txt",
287                "fare_media.txt",
288                "rider_categories.txt",
289                "frequencies.txt",
290                "transfers.txt",
291                "pathways.txt",
292                "levels.txt",
293                "feed_info.txt",
294                "shapes.txt",
295                "translations.txt",
296                "ticketing_deep_links.txt",
297                "ticketing_identifiers.txt",
298                "attributions.txt",
299                "licenses.txt",
300            ] {
301                let path = std::path::Path::new(archive_file.name());
302                if path.file_name() == Some(std::ffi::OsStr::new(gtfs_file)) {
303                    file_mapping.insert(gtfs_file, i);
304                    break;
305                }
306            }
307        }
308
309        let mut result = RawGtfs {
310            agencies: self.read_file(&file_mapping, &mut archive, "agency.txt"),
311            calendar: self.read_optional_file(&file_mapping, &mut archive, "calendar.txt"),
312            calendar_dates: self.read_optional_file(
313                &file_mapping,
314                &mut archive,
315                "calendar_dates.txt",
316            ),
317            routes: self.read_file(&file_mapping, &mut archive, "routes.txt"),
318            stops: self.read_file(&file_mapping, &mut archive, "stops.txt"),
319            stop_times: if self.reader.read_stop_times {
320                self.read_file(&file_mapping, &mut archive, "stop_times.txt")
321            } else {
322                Ok(Vec::new())
323            },
324            trips: self.read_file(&file_mapping, &mut archive, "trips.txt"),
325            fare_attributes: self.read_optional_file(
326                &file_mapping,
327                &mut archive,
328                "fare_attributes.txt",
329            ),
330            fare_rules: self.read_optional_file(&file_mapping, &mut archive, "fare_rules.txt"),
331            fare_products: self.read_optional_file(
332                &file_mapping,
333                &mut archive,
334                "fare_products.txt",
335            ),
336            fare_media: self.read_optional_file(&file_mapping, &mut archive, "fare_media.txt"),
337            rider_categories: self.read_optional_file(
338                &file_mapping,
339                &mut archive,
340                "rider_categories.txt",
341            ),
342            frequencies: self.read_optional_file(&file_mapping, &mut archive, "frequencies.txt"),
343            transfers: self.read_optional_file(&file_mapping, &mut archive, "transfers.txt"),
344            pathways: self.read_optional_file(&file_mapping, &mut archive, "pathways.txt"),
345            levels: self.read_optional_file(&file_mapping, &mut archive, "levels.txt"),
346            feed_info: self.read_optional_file(&file_mapping, &mut archive, "feed_info.txt"),
347            shapes: if self.reader.read_shapes {
348                self.read_optional_file(&file_mapping, &mut archive, "shapes.txt")
349            } else {
350                Some(Ok(Vec::new()))
351            },
352            translations: self.read_optional_file(&file_mapping, &mut archive, "translations.txt"),
353            ticketing_deep_links: self.read_optional_file(
354                &file_mapping,
355                &mut archive,
356                "ticketing_deep_links.txt",
357            ),
358            ticketing_identifiers: self.read_optional_file(
359                &file_mapping,
360                &mut archive,
361                "ticketing_identifiers.txt",
362            ),
363            attributions: self.read_optional_file(&file_mapping, &mut archive, "attributions.txt"),
364            licenses: self.read_optional_file(&file_mapping, &mut archive, "licenses.txt"),
365            read_duration: start_of_read_instant.elapsed(),
366            files,
367            source_format: crate::SourceFormat::Zip,
368            sha256: Some(base16ct::lower::encode_string(&hash)),
369        };
370
371        if self.reader.unkown_enum_as_default {
372            result.unknown_to_default();
373        }
374        Ok(result)
375    }
376
377    fn read_objs<T, O>(&self, mut reader: T, file_name: &str) -> Result<Vec<O>, Error>
378    where
379        for<'de> O: Deserialize<'de>,
380        T: std::io::Read,
381    {
382        let mut bom = [0; 3];
383        reader
384            .read_exact(&mut bom)
385            .map_err(|e| Error::NamedFileIO {
386                file_name: file_name.to_owned(),
387                source: Box::new(e),
388            })?;
389
390        let chained = if bom != [0xefu8, 0xbbu8, 0xbfu8] {
391            bom.chain(reader)
392        } else {
393            [].chain(reader)
394        };
395
396        let mut reader = csv::ReaderBuilder::new()
397            .flexible(true)
398            .trim(if self.reader.trim_fields {
399                csv::Trim::Fields
400            } else {
401                csv::Trim::None
402            })
403            .from_reader(chained);
404        // We store the headers to be able to return them in case of errors
405        let headers = reader
406            .headers()
407            .map_err(|e| Error::CSVError {
408                file_name: file_name.to_owned(),
409                source: e,
410                line_in_error: None,
411            })?
412            .clone()
413            .into_iter()
414            .map(|x| x.trim())
415            .collect::<csv::StringRecord>();
416
417        // Pre-allocate a StringRecord for performance reasons
418        let mut rec = csv::StringRecord::new();
419        let mut objs = Vec::new();
420
421        // Read each record into the pre-allocated StringRecord one at a time
422        while reader.read_record(&mut rec).map_err(|e| Error::CSVError {
423            file_name: file_name.to_owned(),
424            source: e,
425            line_in_error: None,
426        })? {
427            let obj = rec
428                .deserialize(Some(&headers))
429                .map_err(|e| Error::CSVError {
430                    file_name: file_name.to_owned(),
431                    source: e,
432                    line_in_error: Some(crate::error::LineError {
433                        headers: headers.into_iter().map(String::from).collect(),
434                        values: rec.into_iter().map(String::from).collect(),
435                    }),
436                })?;
437            objs.push(obj);
438        }
439        Ok(objs)
440    }
441
442    fn read_objs_from_path<O>(&self, path: std::path::PathBuf) -> Result<Vec<O>, Error>
443    where
444        for<'de> O: Deserialize<'de>,
445    {
446        let file_name = path
447            .file_name()
448            .and_then(|f| f.to_str())
449            .unwrap_or("invalid_file_name")
450            .to_string();
451        if path.exists() {
452            File::open(path)
453                .map_err(|e| Error::NamedFileIO {
454                    file_name: file_name.to_owned(),
455                    source: Box::new(e),
456                })
457                .and_then(|r| self.read_objs(r, &file_name))
458        } else {
459            Err(Error::MissingFile(file_name))
460        }
461    }
462
463    fn read_objs_from_optional_path<O>(
464        &self,
465        dir_path: &std::path::Path,
466        file_name: &str,
467    ) -> Option<Result<Vec<O>, Error>>
468    where
469        for<'de> O: Deserialize<'de>,
470    {
471        File::open(dir_path.join(file_name))
472            .ok()
473            .map(|r| self.read_objs(r, file_name))
474    }
475
476    fn read_file<O, T>(
477        &self,
478        file_mapping: &HashMap<&&str, usize>,
479        archive: &mut zip::ZipArchive<T>,
480        file_name: &str,
481    ) -> Result<Vec<O>, Error>
482    where
483        for<'de> O: Deserialize<'de>,
484        T: std::io::Read + std::io::Seek,
485    {
486        self.read_optional_file(file_mapping, archive, file_name)
487            .unwrap_or_else(|| Err(Error::MissingFile(file_name.to_owned())))
488    }
489
490    fn read_optional_file<O, T>(
491        &self,
492        file_mapping: &HashMap<&&str, usize>,
493        archive: &mut zip::ZipArchive<T>,
494        file_name: &str,
495    ) -> Option<Result<Vec<O>, Error>>
496    where
497        for<'de> O: Deserialize<'de>,
498        T: std::io::Read + std::io::Seek,
499    {
500        file_mapping.get(&file_name).map(|i| {
501            self.read_objs(
502                archive.by_index(*i).map_err(|e| Error::NamedFileIO {
503                    file_name: file_name.to_owned(),
504                    source: Box::new(e),
505                })?,
506                file_name,
507            )
508        })
509    }
510}
gtfs_structures/gtfs_reader.rs

gtfs_structures/
gtfs_reader.rs
