Merge pull request #49 from cauliyang/dev

feat: add func to fq

Cargo.toml

@@ -13,9 +13,9 @@ license = "Apache-2.0"
 
 [workspace.dependencies]
 pyo3 = { version = "0.23.3", features = [
-    "abi3-py39",
-    "extension-module",
-    "anyhow",
+  "abi3-py39",
+  "extension-module",
+  "anyhow",
 ] }
 pyo3-stub-gen = "0.6.2"
 thiserror = "2.0"
@@ -25,13 +25,13 @@ rayon = { version = "1.10" }
 log = "0.4"
 pyo3-log = "0.12.1"
 noodles = { version = "0.87.0", features = [
-    "bgzf",
-    "core",
-    "csi",
-    "fasta",
-    "fastq",
-    "sam",
-    "bam",
+  "bgzf",
+  "core",
+  "csi",
+  "fasta",
+  "fastq",
+  "sam",
+  "bam",
 ] }
 
 bio = "2.0"
@@ -58,7 +58,7 @@ candle-core = { git = "https://github.com/huggingface/candle.git", version = "0.
 colored = "2.2"
 textwrap = "0.16"
 flate2 = { version = "1.0.35", features = [
-    "zlib-ng",
+  "zlib-ng",
 ], default-features = false }
 
 deepbiop-fq = { version = "0.1.12", path = "crates/deepbiop-fq" }

crates/deepbiop-fq/src/python.rs

@@ -1,3 +1,4 @@
+use bstr::BString;
 use numpy::IntoPyArray;
 use std::path::PathBuf;
 
@@ -9,7 +10,7 @@ use crate::{
     utils,
 };
 
-use ahash::HashMap;
+use ahash::{HashMap, HashSet};
 use anyhow::Result;
 use log::warn;
 use needletail::Sequence;
@@ -500,6 +501,21 @@ pub fn fastq_to_fasta(fastq_path: PathBuf, fasta_path: PathBuf) -> Result<()> {
     Ok(())
 }
 
+#[gen_stub_pyfunction(module = "deepbiop.fq")]
+#[pyfunction(name = "select_record_from_fq")]
+pub fn py_select_record_from_fq(
+    selected_reads: Vec<String>,
+    fq: PathBuf,
+    output: PathBuf,
+) -> Result<()> {
+    let selected_reads: HashSet<BString> =
+        selected_reads.into_par_iter().map(|s| s.into()).collect();
+
+    let records = io::select_record_from_fq(fq, &selected_reads)?;
+    io::write_fq_for_noodle_record(&records, output)?;
+    Ok(())
+}
+
 // register fq sub_module
 pub fn register_fq_module(parent_module: &Bound<'_, PyModule>) -> PyResult<()> {
     let sub_module_name = "fq";
@@ -512,6 +528,7 @@ pub fn register_fq_module(parent_module: &Bound<'_, PyModule>) -> PyResult<()> {
     child_module.add_class::<encode::ParquetEncoder>()?;
     child_module.add_class::<predicts::Predict>()?;
 
+    child_module.add_function(wrap_pyfunction!(py_select_record_from_fq, &child_module)?)?;
     child_module.add_function(wrap_pyfunction!(fastq_to_fasta, &child_module)?)?;
     child_module.add_function(wrap_pyfunction!(get_label_region, &child_module)?)?;
     child_module.add_function(wrap_pyfunction!(encode_qual, &child_module)?)?;