diff --git a/otdrs.pyi b/otdrs.pyi index 144cef1..81026e8 100644 --- a/otdrs.pyi +++ b/otdrs.pyi @@ -272,6 +272,14 @@ class SORFile: data_points: DataPoints | None proprietary_blocks: list[ProprietaryBlock] + def to_bytes(self) -> bytes: + """Returns the SOR file as a byte string.""" + ... + + def write_file(self, path: str) -> None: + """Writes the SOR file to the given path.""" + ... + def parse_file(path: str) -> SORFile: """Load a SOR from the given path and parse it""" diff --git a/src/python.rs b/src/python.rs index f5629f6..0d3623a 100644 --- a/src/python.rs +++ b/src/python.rs @@ -4,7 +4,8 @@ use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use pyo3::types::PyBytes; use std::fs::File; -use std::io::Read; +use std::io::{Read, Write}; + /// Loads an OTDR file and returns the result #[pyfunction] fn parse_file(path: String) -> PyResult { @@ -30,10 +31,36 @@ fn parse_bytes(bytes: &Bound<'_, PyBytes>) -> PyResult { return result; } +#[pymethods] +impl SORFile { + /// Returns the SOR file as a byte string. + #[pyo3(name = "to_bytes")] + fn to_bytes_py<'py>(&self, py: Python<'py>) -> PyResult> { + match self.to_bytes() { + Ok(bytes) => Ok(PyBytes::new(py, &bytes)), + Err(err) => Err(PyRuntimeError::new_err(err.to_string())), + } + } + + /// Writes the SOR file to the given path. + #[pyo3(name = "write_file")] + fn write_file_py(&self, path: String) -> PyResult<()> { + match self.to_bytes() { + Ok(bytes) => { + let mut file = std::fs::File::create(path)?; + file.write_all(&bytes)?; + Ok(()) + } + Err(err) => Err(PyRuntimeError::new_err(err.to_string())), + } + } +} + /// This module is implemented in Rust. #[pymodule] fn otdrs(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(parse_file, m)?)?; m.add_function(wrap_pyfunction!(parse_bytes, m)?)?; + m.add_class::()?; return Ok(()); } diff --git a/tests/__pycache__/test_python_otdrs.cpython-311-pytest-8.4.1.pyc b/tests/__pycache__/test_python_otdrs.cpython-311-pytest-8.4.1.pyc new file mode 100644 index 0000000..67cc272 Binary files /dev/null and b/tests/__pycache__/test_python_otdrs.cpython-311-pytest-8.4.1.pyc differ diff --git a/tests/test_python_otdrs.py b/tests/test_python_otdrs.py new file mode 100644 index 0000000..5f10615 --- /dev/null +++ b/tests/test_python_otdrs.py @@ -0,0 +1,43 @@ +import otdrs +import os +import tempfile + +def test_roundtrip(): + """ + Tests that a SOR file can be read, written to bytes, and read back again + without changing the content. + """ + original_sor = otdrs.parse_file("data/example1-noyes-ofl280.sor") + + sor_bytes = original_sor.to_bytes() + roundtrip_sor = otdrs.parse_bytes(sor_bytes) + + # The map block is re-calculated on write, so it will be different. + # We can't compare it directly. + # Let's compare the other blocks. + assert original_sor.general_parameters == roundtrip_sor.general_parameters + assert original_sor.supplier_parameters == roundtrip_sor.supplier_parameters + assert original_sor.fixed_parameters == roundtrip_sor.fixed_parameters + assert original_sor.key_events == roundtrip_sor.key_events + assert original_sor.link_parameters == roundtrip_sor.link_parameters + assert original_sor.data_points == roundtrip_sor.data_points + assert original_sor.proprietary_blocks == roundtrip_sor.proprietary_blocks + + # Test write_file() + fd, temp_filename = tempfile.mkstemp(suffix=".sor") + os.close(fd) + + try: + original_sor.write_file(temp_filename) + roundtrip_sor_from_file = otdrs.parse_file(temp_filename) + + assert original_sor.general_parameters == roundtrip_sor_from_file.general_parameters + assert original_sor.supplier_parameters == roundtrip_sor_from_file.supplier_parameters + assert original_sor.fixed_parameters == roundtrip_sor_from_file.fixed_parameters + assert original_sor.key_events == roundtrip_sor_from_file.key_events + assert original_sor.link_parameters == roundtrip_sor_from_file.link_parameters + assert original_sor.data_points == roundtrip_sor_from_file.data_points + assert original_sor.proprietary_blocks == roundtrip_sor_from_file.proprietary_blocks + + finally: + os.remove(temp_filename)