From 550d6111f383724ecf0344402fe472af40e79bc2 Mon Sep 17 00:00:00 2001 From: Gerard CL Date: Sun, 26 Nov 2023 17:34:53 +0100 Subject: [PATCH] enable station checks - enhanced python binding with use of PyResult and PyErr --- src/cli.rs | 32 ++++++++++++------------- src/lib.rs | 3 +-- src/stations.rs | 62 +++++++++++++++++++++++++++++------------------- src/timetable.rs | 50 ++++++++++++++++++++++---------------- 4 files changed, 84 insertions(+), 63 deletions(-) diff --git a/src/cli.rs b/src/cli.rs index 19feb92..9a2d89b 100644 --- a/src/cli.rs +++ b/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 = 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::() - .unwrap(); + .parse::()?; 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) { diff --git a/src/lib.rs b/src/lib.rs index 70475f2..261dc8a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,7 @@ use pyo3::prelude::*; mod stations; -use stations::{load_stations, Renfe}; +use stations::Renfe; mod timetable; use timetable::{print_timetable, search_timetable}; mod cli; @@ -15,7 +15,6 @@ fn renfe_cli(_py: Python, m: &PyModule) -> PyResult<()> { m.add_class::()?; 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(()) diff --git a/src/stations.rs b/src/stations.rs index fca963c..a3efbc1 100644 --- a/src/stations.rs +++ b/src/stations.rs @@ -1,27 +1,9 @@ -use pyo3::{pyclass, pyfunction, pymethods, PyResult}; +use pyo3::{ + exceptions::{PyConnectionError, PyValueError}, + pyclass, pymethods, PyResult, +}; use scraper::{Html, Selector}; -#[pyfunction] -pub fn load_stations() -> Vec { - 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(_) => { panic!("something wrong") } - }; - - let parsed_html = Html::parse_document(&response.into_string().unwrap()); - - let selector = &Selector::parse(r#"#O > option"#).unwrap(); - - let stations: Vec = parsed_html - .select(selector) - .flat_map(|el| el.text()) - .map(|t| t.to_string()) - .collect(); - - stations[1..].to_vec() -} - #[pyclass] pub struct Renfe { stations: Vec, @@ -31,12 +13,28 @@ pub struct Renfe { impl Renfe { #[new] pub fn new() -> PyResult { + 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 stations: Vec = parsed_html + .select(selector) + .flat_map(|el| el.text()) + .map(|t| t.to_string()) + .collect(); + Ok(Renfe { - stations: load_stations(), + stations: stations[1..].to_vec(), }) } - fn check(&self, station: String) -> PyResult> { + pub fn stations_match(&self, station: String) -> PyResult> { let found: Vec<&String> = self .stations .iter() @@ -44,4 +42,20 @@ impl Renfe { .collect(); Ok(found) } + + pub fn filter_station(&self, station: String) -> PyResult { + 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), + } + } } diff --git a/src/timetable.rs b/src/timetable.rs index 932f7b4..c1983ae 100644 --- a/src/timetable.rs +++ b/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> { +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>> { 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>) { #[cfg(test)] mod tests { + use chrono::{Datelike, Utc}; + use crate::{print_timetable, search_timetable}; #[test] - fn test_search_timetable() -> Result<(), Box> { - // 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> { + 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(()) }