feat: Add SOR file writing to Python library

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"""
 

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<SORFile> {
@@ -30,10 +31,36 @@ fn parse_bytes(bytes: &Bound<'_, PyBytes>) -> PyResult<SORFile> {
     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<Bound<'py, PyBytes>> {
+        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::<SORFile>()?;
     return Ok(());
 }

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)