Huge overhaul

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "shr_parser_py"
 description = "Python bindings for the shr_parser Rust crate"
-version = "0.1.5"
+version = "0.1.6"
 edition = "2021"
 license = "GPL-3.0-only"
 authors = ["Lasse Nielsen <[email protected]>"]

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "shr_parser"
-requires-python = ">=3.8,<=3.12"
+requires-python = ">=3.8,<=3.13"
 classifiers = [
     "Programming Language :: Rust",
     "Programming Language :: Python :: Implementation :: CPython",

src/lib.rs

@@ -1,25 +1,110 @@
+use fmt::Display;
+use std::fmt;
 use pyo3::prelude::*;
 use pyo3::wrap_pyfunction;
 use std::path::PathBuf;
 use ::shr_parser::{SHRParser, SHRParsingType};
 
 /// A wrapper around the SHRParser for Python
-#[pyclass]
+#[pyclass(name="SHRParser", subclass)]
 struct PySHRParser {
     parser: SHRParser,
+    parsing_type: PySHRParsingType,
+}
+
+#[pyclass(name="SHRSweep")]
+struct PySHRSweep {
+    sweep_number: i32,
+    timestamp: u64,
+    frequency: f64,
+    amplitude: f64,
+}
+
+#[pyclass(name="SHRParsingType", eq, eq_int)]
+#[derive(Clone, Copy, Debug, PartialEq)]
+enum PySHRParsingType {
+    PEAK = 0,
+    MEAN = 1,
+    LOW = 2,
+}
+
+impl TryFrom<i32> for PySHRParsingType {
+    type Error = &'static str;
+
+    fn try_from(value: i32) -> Result<Self, Self::Error> {
+        match value {
+            0 => Ok(PySHRParsingType::PEAK),
+            1 => Ok(PySHRParsingType::MEAN),
+            2 => Ok(PySHRParsingType::LOW),
+            _ => Err("Invalid value for SHRParsingType"),
+        }
+    }
+}
+
+impl TryFrom<PySHRParsingType> for SHRParsingType {
+    type Error = &'static str;
+
+    fn try_from(value: PySHRParsingType) -> Result<Self, Self::Error> {
+        match value {
+            PySHRParsingType::PEAK => Ok(SHRParsingType::Peak),
+            PySHRParsingType::MEAN => Ok(SHRParsingType::Mean),
+            PySHRParsingType::LOW => Ok(SHRParsingType::Low),
+        }
+    }
+}
+
+impl Display for PySHRParsingType {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            PySHRParsingType::PEAK => write!(f, "SHRParsingType.PEAK"),
+            PySHRParsingType::MEAN => write!(f, "SHRParsingType.MEAN"),
+            PySHRParsingType::LOW => write!(f, "SHRParsingType.LOW"),
+        }
+    }
+}
+
+#[pymethods]
+impl PySHRSweep {
+    fn __repr__(&self) -> String {
+        format!(
+            "SHRSweep(sweep_number={}, timestamp={}, frequency={}, amplitude={})",
+            self.sweep_number, self.timestamp, self.frequency, self.amplitude
+        )
+    }
+
+    #[getter]
+    fn sweep_number(&self) -> i32 {
+        self.sweep_number
+    }
+
+    #[getter]
+    fn timestamp(&self) -> u64 {
+        self.timestamp
+    }
+
+    #[getter]
+    fn frequency(&self) -> f64 {
+        self.frequency
+    }
+
+    #[getter]
+    fn amplitude(&self) -> f64 {
+        self.amplitude
+    }
 }
 
 #[pymethods]
 impl PySHRParser {
     #[new]
-    fn new(file_path: String, parsing_type: i32) -> PyResult<Self> {
-        let parsing_type = SHRParsingType::try_from(parsing_type).map_err(|e| {
-            pyo3::exceptions::PyValueError::new_err(format!("Invalid parsing type: {}", e))
-        })?;
-        let parser = SHRParser::new(PathBuf::from(file_path), parsing_type).map_err(|e| {
+    fn new(file_path: &str, parsing_type: PySHRParsingType) -> PyResult<Self> {
+        let parser = SHRParser::new(PathBuf::from(file_path), SHRParsingType::try_from(parsing_type).unwrap()).map_err(|e| {
             pyo3::exceptions::PyIOError::new_err(format!("Failed to parse SHR file: {:?}", e))
         })?;
-        Ok(PySHRParser { parser })
+        Ok(PySHRParser { parser, parsing_type })
+    }
+
+    fn __repr__(&self) -> String {
+        format!("SHRParser(file_path='{}', parsing_type={})", self.parser.get_file_path().to_string_lossy(), self.parsing_type)
     }
 
     fn to_csv(&self, path: String) -> PyResult<()> {
@@ -28,11 +113,16 @@ impl PySHRParser {
         })
     }
 
-    fn get_sweeps(&self) -> PyResult<Vec<(i32, u64, f64, f64)>> {
+    fn get_sweeps(&self) -> PyResult<Vec<PySHRSweep>> {
         let sweeps = self.parser.get_sweeps();
         Ok(sweeps
             .into_iter()
-            .map(|sweep| (sweep.sweep_number, sweep.timestamp, sweep.frequency, sweep.amplitude))
+            .map(|sweep| PySHRSweep {
+                sweep_number: sweep.sweep_number,
+                timestamp: sweep.timestamp,
+                frequency: sweep.frequency,
+                amplitude: sweep.amplitude,
+            })
             .collect())
     }
 
@@ -46,16 +136,19 @@ impl PySHRParser {
     }
 }
 
+
 /// Create a new SHRParser instance.
 #[pyfunction]
-fn create_parser(file_path: String, parsing_type: i32) -> PyResult<PySHRParser> {
-    PySHRParser::new(file_path, parsing_type)
+fn create_parser(file_path: &str, parsing_type: i32) -> PyResult<PySHRParser> {
+    PySHRParser::new(file_path, parsing_type.try_into().unwrap())
 }
 
 /// A Python module implemented in Rust.
 #[pymodule]
  fn shr_parser(module: &Bound<'_, PyModule>) -> PyResult<()> {
     module.add_class::<PySHRParser>()?;
+    module.add_class::<PySHRSweep>()?;
+    module.add_class::<PySHRParsingType>()?;
     module.add_function(wrap_pyfunction!(create_parser, module)?)?;
     Ok(())
 }