1use pyo3::prelude::*;
2use pyo3::wrap_pyfunction;
3use std::path::PathBuf;
4use shr_parser::{SHRParser, SHRParsingType};
5
6#[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#[pyfunction]
51fn create_parser(file_path: String, parsing_type: i32) -> PyResult<PySHRParser> {
52 PySHRParser::new(file_path, parsing_type)
53}
54
55#[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}