Add Serialisation Support

Cargo.toml

@@ -14,3 +14,10 @@ sourmash = "0.15.1"
 anyhow = "1.0.89"
 log = "0.4.22"
 env_logger = "0.11.5"
+
+# For JSON serialization/deserialization
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+
+# For Gzip compression/decompression
+flate2 = "1.0"

src/lib.rs

@@ -1,19 +1,29 @@
 // Standard library imports
 use std::collections::hash_map::IntoIter;
 use std::collections::{HashMap, HashSet};
+use std::fs::File;
+use std::io::{Read, Write};
+use std::path::Path;
 
 // External crate imports
 use anyhow::{anyhow, Result};
 use log::debug;
-use pyo3::exceptions::PyValueError;
+use pyo3::exceptions::{PyIOError, PyValueError};
 use pyo3::prelude::*;
 use sourmash::encodings::HashFunctions;
 use sourmash::signature::SeqToHashes;
 
+use flate2::read::GzDecoder;
+use flate2::write::GzEncoder;
+use flate2::Compression;
+use serde::{Deserialize, Serialize};
+use std::io;
+
 // Set version variable
 const VERSION: &str = env!("CARGO_PKG_VERSION");
 
 #[pyclass]
+#[derive(Serialize, Deserialize, Debug)]
 struct KmerCountTable {
     counts: HashMap<u64, u64>,
     pub ksize: u8,
@@ -179,11 +189,59 @@ impl KmerCountTable {
         Ok(to_remove.len() as u64)
     }
 
-    // TODO: Serialize the KmerCountTable instance to a JSON string.
+    // Serialize the KmerCountTable as a JSON string
+    pub fn serialize_json(&self) -> PyResult<String> {
+        serde_json::to_string(&self)
+            .map_err(|e| PyValueError::new_err(format!("Serialization failed: {}", e)))
+    }
+
+    // Save the KmerCountTable to a Gzip-compressed JSON file
+    pub fn save(&self, filepath: String) -> PyResult<()> {
+        let serialized = self.serialize_json()?;
+
+        let path = Path::new(&filepath);
+        let file = File::create(path)
+            .map_err(|e| PyValueError::new_err(format!("File creation failed: {}", e)))?;
+        let mut encoder = GzEncoder::new(file, Compression::default());
+
+        encoder
+            .write_all(serialized.as_bytes())
+            .map_err(|e| PyValueError::new_err(format!("Failed to write to file: {}", e)))?;
+        encoder
+            .finish()
+            .map_err(|e| PyValueError::new_err(format!("Failed to finish compression: {}", e)))?;
+
+        Ok(())
+    }
+
+    #[staticmethod]
+    pub fn load(filepath: String) -> PyResult<Self> {
+        let path = Path::new(&filepath);
+        let file = File::open(path)
+            .map_err(|e| PyValueError::new_err(format!("File open failed: {}", e)))?;
+        let mut decoder = GzDecoder::new(file);
+        let mut decompressed_data = String::new();
 
-    // TODO: Compress JSON string with gzip and save to file
+        decoder
+            .read_to_string(&mut decompressed_data)
+            .map_err(|e| PyValueError::new_err(format!("Decompression failed: {}", e)))?;
 
-    // TODO: Static method to load KmerCountTable from serialized JSON. Yield new object.
+        let loaded_table: KmerCountTable = serde_json::from_str(&decompressed_data)
+            .map_err(|e| PyValueError::new_err(format!("Deserialization failed: {}", e)))?;
+
+        // Check for version mismatch
+        if loaded_table.version != VERSION {
+            eprint!(
+                "Warning: Version mismatch: loaded version is {}, but current version is {}",
+                loaded_table.version, VERSION
+            );
+            io::stderr()
+                .flush()
+                .map_err(|e| PyIOError::new_err(e.to_string()))?; // Flush stderr
+        }
+
+        Ok(loaded_table)
+    }
 
     // TODO: Add method "dump"
     // Output tab delimited kmer:count pairs

src/python/tests/test_serialization.py

@@ -0,0 +1,84 @@
+import gzip
+import json
+import tempfile
+import pytest
+
+from oxli import KmerCountTable
+from os import remove
+from test_attr import get_version_from_cargo_toml
+
+CURRENT_VERSION = get_version_from_cargo_toml()
+
+@pytest.fixture
+def sample_kmer_table():
+    """Fixture that provides a sample KmerCountTable object."""
+    table = KmerCountTable(ksize=4)
+    table.count("AAAA")
+    table.count("TTTT")
+    return table
+
+@pytest.fixture
+def temp_file():
+    """Fixture that provides a temporary file path for testing."""
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.json.gz') as temp:
+        yield temp.name
+    # Remove the file after the test is done
+    remove(temp.name)
+
+def test_serialize_json(sample_kmer_table):
+    """
+    Test case for the `serialize_json` function.
+
+    This test verifies that the `serialize_json` function correctly serializes a
+    KmerCountTable object into a JSON string.
+    """
+    # Serialize the KmerCountTable object to JSON
+    json_data = sample_kmer_table.serialize_json()
+
+    # Convert back to dict to verify correctness
+    json_dict = json.loads(json_data)
+
+    # Check that essential attributes exist
+    assert "counts" in json_dict, "Counts should be serialized."
+    assert json_dict["ksize"] == 4, "Ksize should be correctly serialized."
+    assert sample_kmer_table.version == json_dict["version"], "Version should be serialized."
+
+def test_save_load_roundtrip(sample_kmer_table, temp_file):
+    """
+    Test the save and load functionality.
+
+    This test saves a KmerCountTable object to a file, then loads it back and
+    verifies that the data in the loaded object matches the original.
+    """
+    # Save the sample KmerCountTable to a Gzip file
+    sample_kmer_table.save(temp_file)
+
+    # Load the KmerCountTable from the file
+    loaded_table = KmerCountTable.load(temp_file)
+
+    # Verify that the loaded data matches the original
+    assert loaded_table.get("AAAA") == sample_kmer_table.get("AAAA"), "Counts should be preserved after loading."
+    assert loaded_table.get("TTTT") == sample_kmer_table.get("TTTT"), "Counts for reverse complement should be preserved."
+    assert list(loaded_table) == list(sample_kmer_table), "All records in same order."
+
+def test_version_warning_on_load_stderr(sample_kmer_table, temp_file, capfd):
+    """
+    Test that a warning is issued if the loaded object's version is different from the current Oxli version.
+
+    Uses pytest's capsys fixture to capture stderr output.
+    """
+    # Save the table to a file
+    sample_kmer_table.save(temp_file)
+
+    # Mock the current version to simulate a version mismatch
+    mock_json = sample_kmer_table.serialize_json().replace(CURRENT_VERSION, "0.0.1")
+    with gzip.open(temp_file, 'wt') as f:
+        json.dump(json.loads(mock_json), f)
+
+    # Capture stderr output
+    loaded_table = KmerCountTable.load(temp_file)
+    captured = capfd.readouterr()
+
+    # Check stderr for the version mismatch warning
+    assert "Version mismatch" in captured.err
+    assert f"loaded version is 0.0.1, but current version is {CURRENT_VERSION}" in captured.err