Merge pull request #183 from gerardcl/feature/check-stations

check stations

src/cli.rs

@@ -1,51 +1,49 @@
 use chrono::{Datelike, Utc};
 use getopts::Options;
-use pyo3::pyfunction;
+use pyo3::{exceptions::PyValueError, pyfunction, PyResult};
 use std::env;
 
-use crate::timetable::{print_timetable, search_timetable};
+use crate::{
+    stations::Renfe,
+    timetable::{print_timetable, search_timetable},
+};
 
 #[pyfunction]
-pub fn main() {
+pub fn main() -> PyResult<()> {
     let args: Vec<String> = env::args().collect();
     let program = args[0].clone();
     let now = Utc::now();
     let opts = set_opts();
+    let renfe = Renfe::new()?;
 
     let matches = match opts.parse(&args[1..]) {
         Ok(m) => m,
         Err(f) => {
-            panic!("{}", f.to_string())
+            return Err(PyValueError::new_err(f.to_string()));
         }
     };
+
     if matches.opt_present("h") {
         print_usage(&program, opts);
-        return;
+        return Ok(());
     }
-    let origin = matches.opt_str("f");
-    let destination = matches.opt_str("t");
+    let origin = renfe.filter_station(matches.opt_str("f").unwrap_or("".to_owned()))?;
+    let destination = renfe.filter_station(matches.opt_str("t").unwrap_or("".to_owned()))?;
     let day = matches.opt_str("d").unwrap_or(now.day().to_string());
     let month = matches.opt_str("m").unwrap_or(now.month().to_string());
     let year = matches.opt_str("y").unwrap_or(now.year().to_string());
     let wait = matches
         .opt_str("w")
         .unwrap_or(2.to_string())
-        .parse::<u64>()
-        .unwrap();
+        .parse::<u64>()?;
 
     println!("Today is: {}-{}-{}", now.year(), now.month(), now.day());
     println!("Searching timetable for date: {}-{}-{}", year, month, day);
 
-    let timetable = search_timetable(
-        origin.unwrap(),
-        destination.unwrap(),
-        day,
-        month,
-        year,
-        wait,
-    );
+    let timetable = search_timetable(origin, destination, day, month, year, wait)?;
 
     print_timetable(timetable);
+    Ok(())
 }
 
 fn print_usage(program: &str, opts: Options) {

src/lib.rs

@@ -1,7 +1,7 @@
 use pyo3::prelude::*;
 
 mod stations;
-use stations::load_stations;
+use stations::Renfe;
 mod timetable;
 use timetable::{print_timetable, search_timetable};
 mod cli;
@@ -12,9 +12,9 @@ use cli::main;
 /// import the module.
 #[pymodule]
 fn renfe_cli(_py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_class::<Renfe>()?;
     m.add_function(wrap_pyfunction!(search_timetable, m)?)?;
     m.add_function(wrap_pyfunction!(print_timetable, m)?)?;
-    m.add_function(wrap_pyfunction!(load_stations, m)?)?;
     m.add_function(wrap_pyfunction!(main, m)?)?;
 
     Ok(())

src/stations.rs

@@ -1,22 +1,61 @@
-use pyo3::pyfunction;
+use pyo3::{
+    exceptions::{PyConnectionError, PyValueError},
+    pyclass, pymethods, PyResult,
+};
 use scraper::{Html, Selector};
 
-#[pyfunction]
-pub fn load_stations() -> Vec<String> {
-    let response = match ureq::get("https://www.renfe.com/content/renfe/es/en/viajar/informacion-util/horarios/app-horarios.html").call() {
-        Ok(response) => { response },
-        Err(_) => { panic!("something wrong") }
-    };
+#[pyclass]
+pub struct Renfe {
+    stations: Vec<String>,
+}
+
+#[pymethods]
+impl Renfe {
+    #[new]
+    pub fn new() -> PyResult<Self> {
+        println!("Loading stations from Renfe web");
+        let response = match ureq::get("https://www.renfe.com/content/renfe/es/en/viajar/informacion-util/horarios/app-horarios.html").call() {
+            Ok(response) => { response },
+            Err(_) => { return Err(PyConnectionError::new_err("something wrong")) }
+        };
+
+        let parsed_html = Html::parse_document(&response.into_string().unwrap());
+
+        let selector = &Selector::parse(r#"#O > option"#).unwrap();
 
-    let parsed_html = Html::parse_document(&response.into_string().unwrap());
+        let stations: Vec<String> = parsed_html
+            .select(selector)
+            .flat_map(|el| el.text())
+            .map(|t| t.to_string())
+            .collect();
 
-    let selector = &Selector::parse(r#"#O > option"#).unwrap();
+        Ok(Renfe {
+            stations: stations[1..].to_vec(),
+        })
+    }
 
-    let stations: Vec<String> = parsed_html
-        .select(selector)
-        .flat_map(|el| el.text())
-        .map(|t| t.to_string())
-        .collect();
+    pub fn stations_match(&self, station: String) -> PyResult<Vec<&String>> {
+        let found: Vec<&String> = self
+            .stations
+            .iter()
+            .filter(|s| s.contains(&station))
+            .collect();
+        Ok(found)
+    }
 
-    stations[1..].to_vec()
+    pub fn filter_station(&self, station: String) -> PyResult<String> {
+        match self.stations_match(station.clone()) {
+            Ok(v) if v.len() == 1 => {
+                println!(
+                    "Provided input '{}' station matches with '{}'...continue",
+                    station, v[0]
+                );
+                Ok(v[0].to_owned())
+            }
+            Ok(v) => Err(PyValueError::new_err(format!(
+                "Provided input '{station}' station does not match one '{v:?}'"
+            ))),
+            Err(e) => Err(e),
+        }
+    }
 }

src/timetable.rs

@@ -1,5 +1,5 @@
 use headless_chrome::{Browser, LaunchOptions};
-use pyo3::pyfunction;
+use pyo3::{pyfunction, PyResult};
 use scraper::{ElementRef, Html, Selector};
 use std::{collections::HashMap, thread::sleep, time::Duration};
 
@@ -23,15 +23,7 @@ fn make_selector(selector: &str) -> Selector {
     Selector::parse(selector).unwrap()
 }
 
-#[pyfunction]
-pub fn search_timetable(
-    origin: String,
-    destination: String,
-    day: String,
-    month: String,
-    year: String,
-    wait: u64,
-) -> Vec<Vec<String>> {
+fn to_renfe_month(month: String) -> String {
     let months: HashMap<&str, &str> = HashMap::from([
         ("1", "Ene"),
         ("2", "Feb"),
@@ -46,7 +38,18 @@ pub fn search_timetable(
         ("11", "Nov"),
         ("12", "Dec"),
     ]);
+    months[month.as_str()].to_owned()
+}
 
+#[pyfunction]
+pub fn search_timetable(
+    origin: String,
+    destination: String,
+    day: String,
+    month: String,
+    year: String,
+    wait: u64,
+) -> PyResult<Vec<Vec<String>>> {
     println!("loading headless chrome browser");
     let browser = Browser::new(LaunchOptions {
         headless: true,
@@ -72,6 +75,8 @@ pub fn search_timetable(
 
     println!("navigating to renfe timetable search page");
     tab.navigate_to("https://www.renfe.com/es/es/viajar/informacion-util/horarios")
+        .unwrap()
+        .wait_until_navigated()
         .unwrap();
 
     // let _jpeg_data = tab.capture_screenshot(
@@ -117,7 +122,7 @@ pub fn search_timetable(
         .unwrap()
         .click()
         .unwrap();
-    tab.type_str(months[&month.as_str()])
+    tab.type_str(&to_renfe_month(month))
         .unwrap()
         .press_key("Enter")
         .unwrap();
@@ -178,7 +183,7 @@ pub fn search_timetable(
         tracks.push(row);
     }
 
-    tracks
+    Ok(tracks)
 }
 
 #[pyfunction]
@@ -200,17 +205,22 @@ pub fn print_timetable(tracks: Vec<Vec<String>>) {
 
 #[cfg(test)]
 mod tests {
+    use chrono::{Datelike, Utc};
+
     use crate::{print_timetable, search_timetable};
 
     #[test]
-    fn test_search_timetable() -> Result<(), Box<dyn std::error::Error>> {
-        // print_timetable(search_timetable(
-        //     "Girona".to_owned(),
-        //     "Barcelona".to_owned(),
-        //     "28".to_owned(),
-        //     "11".to_owned(),
-        //     "2023".to_owned(),
-        // ));
+    fn test_search_and_print_timetable() -> Result<(), Box<dyn std::error::Error>> {
+        let now = Utc::now();
+
+        print_timetable(search_timetable(
+            "Girona".to_owned(),
+            "Barcelona".to_owned(),
+            now.day().to_string(),
+            now.month().to_string(),
+            now.year().to_string(),
+            2,
+        )?);
 
         Ok(())
     }