properly handle (and test) errors

oxli-bio · Sep 2, 2024 · c8fe026 · c8fe026
1 parent 40381ad
commit c8fe026
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 5 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -85,11 +85,12 @@ impl KmerCountTable {
     }
 
     // Consume this DNA strnig. Return number of k-mers consumed.
-    pub fn consume(&mut self, seq: String) -> PyResult<u64> {
+    #[pyo3(signature = (seq, allow_bad_kmers=true))]
+    pub fn consume(&mut self, seq: String, allow_bad_kmers: bool) -> PyResult<u64> {
         let hashes = SeqToHashes::new(
             seq.as_bytes(),
             self.ksize.into(),
-            false,
+            allow_bad_kmers,
             false,
             HashFunctions::Murmur64Dna,
             42,
@@ -104,8 +105,12 @@ impl KmerCountTable {
                     self.count_hash(x);
                     ()
                 }
-                Err(err) => break, // @CTB
+                Err(_) => {
+                    let msg = format!("bad k-mer encountered at position {}", n);
+                    return Err(PyValueError::new_err(msg));
+                }
             }
+
             n += 1;
         }
 

diff --git a/src/python/tests/test_basic.py b/src/python/tests/test_basic.py
@@ -38,8 +38,27 @@ def test_consume_2():
     assert cg.get("TCGG") == 1
     assert cg.get("CCGA") == 1 # reverse complement!
 
+
 def test_consume_bad_DNA():
     cg = oxli.KmerCountTable(4)
     seq = "ATCGGX"
-    print(cg.consume(seq))
-    assert 0
+    with pytest.raises(ValueError,
+                       match="bad k-mer encountered at position 2"):
+        cg.consume(seq, allow_bad_kmers=False)
+
+
+def test_consume_bad_DNA_2():
+    cg = oxli.KmerCountTable(4)
+    seq = "XATCGG"
+    with pytest.raises(ValueError,
+                       match="bad k-mer encountered at position 0"):
+        cg.consume(seq, allow_bad_kmers=False)
+
+
+def test_consume_bad_DNA_ignore():
+    cg = oxli.KmerCountTable(4)
+    seq = "XATCGG"
+    print(cg.consume(seq, allow_bad_kmers=True))
+    assert cg.get("ATCG") == 1
+    assert cg.get("TCGG") == 1
+    assert cg.get("CCGA") == 1 # rc