Repository: Adoni5/mappy-rs (language: python)

Issue:
Allow map batch to accept lists
It would be nice if `Aligner.map_batch` could accept a list of input, currently:

```python
import mappy_rs
fasta = [
  {'id': 0, 'seq': 'AGAGTGAAGCCAATATTCCGATAACGATTGCTTTCATGATATCCCTCATTCATCACAAGTTTT'},
  {'id': 1, 'seq': 'AGAGTGAAGCCAATATTCCGATAACGATTGCTTTCATGATATCCCTCATTCATCACAAGTTTT'},
]
al = mappy_rs.Aligner("resources/test/test.mmi")
al.enable_threading(2)
mappings = al.map_batch(fasta)
# TypeError: argument 'seqs': 'list' object cannot be converted to 'Iterator'
```


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 3acffb4..f275fdf 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -3,6 +3,7 @@ name: CI
 on:
   create:
   pull_request:
+  workflow_dispatch:
 
 jobs:
   linux:
diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml
index 34034fd..2fc7a8c 100644
--- a/.github/workflows/check.yml
+++ b/.github/workflows/check.yml
@@ -3,6 +3,7 @@ on:
   push:
     branches: [main]
   pull_request:
+  workflow_dispatch:
 name: check
 jobs:
   fmt:
diff --git a/README.md b/README.md
index 884370c..748f193 100644
--- a/README.md
+++ b/README.md
@@ -42,28 +42,30 @@ aligner = mappy_rs.Aligner("resources/test/test.mmi")
 
 The current iteration of `mappy-rs` serves as a drop in for `mappy`, implementing all the same methods. However if this is the use case, you may well be better off using `mappy`, as the extra level of Rust betwene your python and C++ may well add slighly slower performance.
 
-#### Multithreading
+### Multithreading
 In order to use multi threading, one must first enable it.
 
 ```python
 import mappy_rs
-aligner = mappy_rs.Aligner("/path/to/index.mmi")
+aligner = mappy_rs.Aligner("resources/test/test.mmi")
 # Use 10 threads
 aligner.enable_threading(10)
 ```
 
-Enabling threading makes the `map_batch` method available. ⚠️ _Currently, this only takes an iterator of dictionaries. The dictionaries can have any number of keys and depth, but __must__ contain one key value pair of "seq": "sequence"_.
+Enabling threading makes the `map_batch` method available.  
+This method requires a list or iterable of dictionaries, which can have any number of keys and depth, but **must** contain the key `seq` with a string value in the top-level dictionary.
 
-For example
+For example:
 
 ```python
 import mappy_rs
 aligner = mappy_rs.Aligner("resources/test/test.mmi")
 aligner.enable_threading(10)
-for (mapping, data) in aligner.map_batch(iter([{"seq": "ACGTAGCATCGAGACTACGA", "Other_random_key": "banter"}, {"seq": "ACGTAGCATCGAGACTACGA", "Other_random_key": "banter"}])):
+seqs = [
+    {"seq": "ACGTAGCATCGAGACTACGA", "Other_random_key": "banter"}, 
+    {"seq": "ACGTAGCATCGAGACTACGA", "Other_random_key": "banter"},
+]
+for (mapping, data) in aligner.map_batch(seqs):
     print(list(mapping))
     print(data)
 ```
-
-Will work, but without the call to `iter()` it would crash.
-
diff --git a/src/lib.rs b/src/lib.rs
index 0f51b3a..e6eb231 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -6,10 +6,13 @@
 
 use crossbeam::channel::{bounded, Receiver, RecvError, Sender};
 use fnv::FnvHashMap;
-use pyo3::exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError};
+use pyo3::exceptions::{
+    PyKeyError, PyNotImplementedError, PyRuntimeError, PyTypeError, PyValueError,
+};
 use pyo3::prelude::*;
 use pyo3::pyclass::IterNextOutput;
-use pyo3::types::{PyIterator, PyTuple};
+use pyo3::types::{PyIterator, PyList, PySequence, PyTuple};
+use pyo3::FromPyObject;
 use std::collections::HashMap;
 use std::fmt::{Display, Formatter};
 use threadpool::ThreadPool;
@@ -494,7 +497,7 @@ impl Aligner {
     ///
     /// Example
     /// -------
-    /// ```aligner::enable_threading(8)```
+    /// `aligner::enable_threading(8)`
     #[pyo3(signature = (n_threads), text_signature = "(n_threads=8)")]
     fn enable_threading(&mut self, n_threads: usize) -> PyResult<()> {
         self.n_threads = n_threads;
@@ -502,7 +505,7 @@ impl Aligner {
     }
 
     /// Align a sequence with optional Metadata tuple. If provided, the tuple MUST be in the shape of (channel_number: int, read_id: str)
-    fn map_batch(&self, seqs: &PyIterator) -> PyResult<AlignmentBatchResultIter> {
+    fn map_batch(&self, seqs: &PyAny) -> PyResult<AlignmentBatchResultIter> {
         let mut res = AlignmentBatchResultIter::new();
         // do the heavy work
         self._map_batch(&mut res, seqs)?;
@@ -602,26 +605,53 @@ impl Aligner {
 
     /// Align a batch of reads provided in an iterator, using a threadpool with the number of threads specified by
     /// .enable_threading()
-    pub fn _map_batch(
-        &self,
-        res: &mut AlignmentBatchResultIter,
-        seqs: &PyIterator,
-    ) -> PyResult<()> {
+    pub fn _map_batch(&self, res: &mut AlignmentBatchResultIter, seqs: &PyAny) -> PyResult<()> {
         if self.n_threads == 0_usize {
             return Err(PyRuntimeError::new_err(
                 "Multi threading not enabled on this instance. Please call `.enable_threading()`",
             ));
         }
+        let _ = match seqs.extract() {
+            Ok(SupportedTypes::List(_)) => (),
+            Ok(SupportedTypes::Tuple(_)) => (),
+            Ok(SupportedTypes::Iter(_)) => (),
+            Ok(SupportedTypes::Sequence(_)) => (),
+            _ => {
+                return Err(PyTypeError::new_err(
+                    "Unsupported batch type, pass a list, iter, generator or tuple",
+                ))
+            }
+        };
+
+        let iter = match seqs.iter() {
+            Ok(it) => it,
+            _ => return Err(PyTypeError::new_err("Could not iterate batch")),
+        };
+
         let p = ThreadPool::new(self.n_threads);
-        for (id_num, py_dicts) in seqs.iter()?.enumerate() {
+        for (id_num, py_dicts) in iter.enumerate() {
             let py_dict = py_dicts?;
-            let data: HashMap<String, Py<PyAny>> = py_dict.extract()?;
+            let data: HashMap<String, Py<PyAny>> = match py_dict.extract() {
+                Ok(x) => x,
+                _ => {
+                    return Err(PyTypeError::new_err(
+                        "Element in iterable is not a dictionary",
+                    ))
+                }
+            };
             res.data.insert(id_num, data);
             let sendy = res.tx.clone();
-            let seq: String = py_dict
-                .get_item("seq")
-                .expect("AHHH Key 🗝️  not found in iterated dictionary")
-                .extract()?;
+            let seq: String = match py_dict.get_item("seq") {
+                Ok(seq) => match seq.extract() {
+                    Ok(seq) => seq,
+                    _ => return Err(PyValueError::new_err("`seq` must be a string")),
+                },
+                _ => {
+                    return Err(PyKeyError::new_err(
+                        "AHHH Key 🗝️  not found in iterated dictionary",
+                    ))
+                }
+            };
             let aligner = self.clone();
             p.execute(move || {
                 let maps = aligner.map(seq, None, true, true).unwrap();
@@ -639,6 +669,19 @@ impl Aligner {
     }
 }
 
+/// Python iterable types that are accepted by the `Aligner.map_batch()` function
+#[derive(FromPyObject)]
+enum SupportedTypes<'py> {
+    /// PyList wrapper
+    List(&'py PyList),
+    /// PyTuple wrapper
+    Tuple(&'py PyTuple),
+    /// PyIter wrapper
+    Iter(&'py PyIterator),
+    /// PySequence wrapper
+    Sequence(&'py PySequence),
+}
+
 /// Struct for returning data to the python runtime as an iterabled.
 #[pyclass]
 pub struct AlignmentBatchResultIter {


Test patch:
diff --git a/tests/python_test.py b/tests/python_test.py
index 5437b5e..fada20c 100644
--- a/tests/python_test.py
+++ b/tests/python_test.py
@@ -41,13 +41,18 @@ def al(mmi_file):
     return mappy_rs.Aligner(mmi_file)
 
 
+@pytest.fixture
+def fasta(request):
+    return request.getfixturevalue(request.param)
+
+
 @pytest.fixture
 def fasta_file():
     return str(FA_FILE)
 
 
 @pytest.fixture
-def fasta_iter(fasta_file):
+def fasta_list(fasta_file):
     with open(fasta_file, "rt") as fh:
         seqs = [s for _, s in read_fasta(fh)]
 
@@ -55,7 +60,22 @@ def fasta_iter(fasta_file):
         {"id": i, "seq": seq}
         for i, seq in enumerate(copy.copy(s) for _ in range(10) for s in seqs)
     ]
-    return iter(seqs)
+    return seqs
+
+
+@pytest.fixture
+def fasta_iter(fasta_list):
+    return iter(fasta_list)
+
+
+@pytest.fixture
+def fasta_tuple(fasta_list):
+    return tuple(fasta_list)
+
+
+@pytest.fixture
+def fasta_generator(fasta_list):
+    return (item for item in fasta_list)
 
 
 def test_test(al):
@@ -114,10 +134,64 @@ def test_map_one(al):
     assert mapping.target_end == 400
 
 
-def test_map_batch(al, fasta_iter):
+@pytest.mark.parametrize(
+    "fasta",
+    ["fasta_iter", "fasta_list", "fasta_tuple", "fasta_generator"],
+    indirect=True,
+)
+def test_map_batch(al, fasta):
     al.enable_threading(2)
-    mappings = al.map_batch(fasta_iter)
+    mappings = al.map_batch(fasta)
     n = 0
     for res in mappings:
         n += 1
     assert n == 40
+
+
+def test_map_batch_fail_dict_single(al, fasta_iter):
+    fasta = next(fasta_iter)
+    al.enable_threading(2)
+    with pytest.raises(TypeError) as excinfo:
+        mappings = al.map_batch(fasta)
+    err = str(excinfo)
+    assert "Unsupported batch type, pass a list, iter, generator or tuple" in err
+
+
+def test_map_batch_fail_dict_many(al, fasta_iter):
+    fasta = {i: dct for i, dct in enumerate(fasta_iter)}
+    al.enable_threading(2)
+    with pytest.raises(TypeError) as excinfo:
+        mappings = al.map_batch(fasta)
+    err = str(excinfo)
+    assert "Unsupported batch type, pass a list, iter, generator or tuple" in err
+
+
+def test_map_batch_fail_list_str(al, fasta_iter):
+    fasta = [dct["seq"] for dct in fasta_iter]
+    al.enable_threading(2)
+    with pytest.raises(TypeError) as excinfo:
+        mappings = al.map_batch(fasta)
+    assert "Element in iterable is not a dictionary" in str(excinfo.value)
+
+
+def test_map_batch_fail_no_seq_key(al, fasta_iter):
+    fasta = [{"SEQ": dct["seq"]} for dct in fasta_iter]
+    al.enable_threading(2)
+    with pytest.raises(KeyError) as excinfo:
+        mappings = al.map_batch(fasta)
+    assert "AHHH Key 🗝️  not found in iterated dictionary" in str(excinfo)
+
+
+def test_map_batch_fail_seq_not_str(al, fasta_iter):
+    fasta = [{"seq": dct["seq"].encode()} for dct in fasta_iter]
+    al.enable_threading(2)
+    with pytest.raises(ValueError) as excinfo:
+        mappings = al.map_batch(fasta)
+    assert "`seq` must be a string" in str(excinfo)
+
+
+def test_map_batch_fail_exhausted_iter(al, fasta_iter):
+    _ = list(fasta_iter)
+    al.enable_threading(2)
+    mappings = al.map_batch(fasta_iter)
+    assert len(list(mappings)) == 0