shr_parser/
lib.rs

1use pyo3::prelude::*;
2use pyo3::wrap_pyfunction;
3use std::path::PathBuf;
4use shr_parser::{SHRParser, SHRParsingType};
5
6/// A wrapper around the SHRParser for Python
7#[pyclass]
8struct PySHRParser {
9    parser: SHRParser,
10}
11
12#[pymethods]
13impl PySHRParser {
14    #[new]
15    fn new(file_path: String, parsing_type: i32) -> PyResult<Self> {
16        let parsing_type = SHRParsingType::try_from(parsing_type).map_err(|e| {
17            pyo3::exceptions::PyValueError::new_err(format!("Invalid parsing type: {}", e))
18        })?;
19        let parser = SHRParser::new(PathBuf::from(file_path), parsing_type).map_err(|e| {
20            pyo3::exceptions::PyIOError::new_err(format!("Failed to parse SHR file: {:?}", e))
21        })?;
22        Ok(PySHRParser { parser })
23    }
24
25    fn to_csv(&self, path: String) -> PyResult<()> {
26        self.parser.to_csv(PathBuf::from(path)).map_err(|e| {
27            pyo3::exceptions::PyIOError::new_err(format!("Failed to write to CSV: {:?}", e))
28        })
29    }
30
31    fn get_sweeps(&self) -> PyResult<Vec<(i32, u64, f64, f64)>> {
32        let sweeps = self.parser.get_sweeps();
33        Ok(sweeps
34            .into_iter()
35            .map(|sweep| (sweep.sweep_number, sweep.timestamp, sweep.frequency, sweep.amplitude))
36            .collect())
37    }
38
39    fn get_file_header(&self) -> PyResult<String> {
40        let header = self.parser.get_file_header();
41        Ok(format!("{:?}", header))
42    }
43
44    fn get_file_path(&self) -> PyResult<String> {
45        Ok(self.parser.get_file_path().to_string_lossy().to_string())
46    }
47}
48
49/// Create a new SHRParser instance.
50#[pyfunction]
51fn create_parser(file_path: String, parsing_type: i32) -> PyResult<PySHRParser> {
52    PySHRParser::new(file_path, parsing_type)
53}
54
55/// A Python module implemented in Rust.
56#[pymodule]
57 fn my_module(module: &Bound<'_, PyModule>) -> PyResult<()> {
58    module.add_class::<PySHRParser>()?;
59    module.add_function(wrap_pyfunction!(create_parser, module)?)?;
60    Ok(())
61}