Merge pull request #19 from NREL/fix/from_csv_file

Fix `from_csv_file` function
NREL · Nov 28, 2023 · 89c2d56 · 89c2d56
2 parents 21dc75d + 66d1d89
commit 89c2d56
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 20 deletions.
diff --git a/python/altrios/demos/set_speed_train_sim_demo.py b/python/altrios/demos/set_speed_train_sim_demo.py
@@ -78,15 +78,8 @@
     alt.resources_root() / "demo_data/link_points.csv")["link points"].tolist()
 link_path = [alt.LinkIdx(int(lp)) for lp in link_points]
 
-# TODO: uncomment and fix
-# speed_trace = alt.SpeedTrace.from_csv_file(
-#     str(alt.resources_root() / "demo_data/speed_trace.csv")
-# )
-df_speed_trace = pd.read_csv(alt.resources_root() / "demo_data/speed_trace.csv")
-speed_trace = alt.SpeedTrace(
-    df_speed_trace['time_seconds'],
-    df_speed_trace['speed_meters_per_second'],
-    None,
+speed_trace = alt.SpeedTrace.from_csv_file(
+    alt.resources_root() / "demo_data/speed_trace.csv"
 )
 
 train_sim: alt.SetSpeedTrainSim = tsb.make_set_speed_train_sim(

diff --git a/rust/altrios-core/src/train/set_speed_train_sim.rs b/rust/altrios-core/src/train/set_speed_train_sim.rs
@@ -12,8 +12,8 @@ use super::train_imports::*;
 
     #[staticmethod]
     #[pyo3(name = "from_csv_file")]
-    fn from_csv_file_py(pathstr: String) -> anyhow::Result<Self> {
-        Self::from_csv_file(&pathstr)
+    fn from_csv_file_py(filepath: &PyAny) -> anyhow::Result<Self> {
+        Self::from_csv_file(PathBuf::extract(filepath)?)
     }
 
     fn __len__(&self) -> usize {
@@ -71,7 +71,7 @@ impl SpeedTrace {
 
     /// method to prevent rust-analyzer from complaining
     pub fn is_empty(&self) -> bool {
-        true // not really possible to create an empty SpeedTrace
+        self.time.is_empty() && self.speed.is_empty() && self.engine_on.is_none()
     }
 
     pub fn push(&mut self, speed_element: SpeedTraceElement) -> anyhow::Result<()> {
@@ -99,25 +99,26 @@ impl SpeedTrace {
     }
 
     /// Load speed trace from csv file
-    pub fn from_csv_file(pathstr: &str) -> Result<Self, anyhow::Error> {
-        let pathbuf = PathBuf::from(&pathstr);
+    pub fn from_csv_file<P: AsRef<Path>>(filepath: P) -> anyhow::Result<Self> {
+        let filepath = filepath.as_ref();
 
         // create empty SpeedTrace to be populated
         let mut st = Self::empty();
 
-        let file = File::open(pathbuf)?;
+        let file = File::open(filepath)?;
         let mut rdr = csv::ReaderBuilder::new()
             .has_headers(true)
             .from_reader(file);
         for result in rdr.deserialize() {
             let st_elem: SpeedTraceElement = result?;
             st.push(st_elem)?;
         }
-        if st.is_empty() {
-            bail!("Invalid SpeedTrace file; SpeedTrace is empty")
-        } else {
-            Ok(st)
-        }
+        ensure!(
+            !st.is_empty(),
+            "Invalid SpeedTrace file {:?}; SpeedTrace is empty",
+            filepath
+        );
+        Ok(st)
     }
 }