Merge pull request #23 from nh13/feat/read_record_set_limit

feat: add method to read a maximum # of records

Author: markschl

src/fasta.rs

@@ -254,6 +254,19 @@ where
     // TODO: in next major version, return Result<bool> instead!
     #[inline]
     pub fn read_record_set(&mut self, rset: &mut RecordSet) -> Option<Result<(), Error>> {
+        self.read_record_set_limited(rset, usize::MAX)
+    }
+
+    /// Updates a [RecordSet](struct.RecordSet.html) with new data. The contents of the internal
+    /// buffer are just copied over to the record set and the positions of all records are found.
+    /// Old data will be erased. Returns `None` if the input reached its end.
+    // TODO: in next major version, return Result<bool> instead!
+    #[inline]
+    pub fn read_record_set_limited(
+        &mut self,
+        rset: &mut RecordSet,
+        max_records: usize,
+    ) -> Option<Result<(), Error>> {
         // after read_record_set(), the state is always Positioned, Parsing or Finished
         match self.state {
             State::New => {
@@ -312,6 +325,9 @@ where
             }
             rset.npos += 1;
             self.increment_record();
+            if rset.npos >= max_records {
+                break;
+            }
         }
 
         rset.buffer.clear();

src/fastq.rs

@@ -211,6 +211,19 @@ where
     // TODO: in next major version, return Result<bool> instead!
     #[inline]
     pub fn read_record_set(&mut self, rset: &mut RecordSet) -> Option<Result<(), Error>> {
+        self.read_record_set_limited(rset, usize::MAX)
+    }
+
+    /// Updates a [RecordSet](struct.RecordSet.html) with new data. The contents of the internal
+    /// buffer are just copied over to the record set and the positions of all records are found.
+    /// Old data will be erased. Returns `None` if the input reached its end.
+    // TODO: in next major version, return Result<bool> instead!
+    #[inline]
+    pub fn read_record_set_limited(
+        &mut self,
+        rset: &mut RecordSet,
+        max_records: usize,
+    ) -> Option<Result<(), Error>> {
         // after read_record_set(), the state is always Positioned, Parsing or Finished
         match self.state {
             State::New => {
@@ -259,6 +272,9 @@ where
             }
             rset.buf_positions.push(self.buf_pos.clone());
             self.increment_record();
+            if rset.buf_positions.len() >= max_records {
+                break;
+            }
         }
 
         rset.buffer.clear();

tests/fasta.rs

@@ -97,6 +97,34 @@ fn test_fasta_seq_lines() {
     assert_eq!(lines, vec![b"AT", b"GC"]);
 }
 
+#[test]
+fn test_fastq_read_record_set_limited() {
+    for max_records in 3..10 {
+        use std::io::Write;
+        let mut fasta_vec = Vec::with_capacity(13 * max_records);
+        for i in 0..max_records {
+            write!(&mut fasta_vec, ">id{i}\nATGC\n").unwrap();
+        }
+
+        let mut reader = Reader::new(&fasta_vec[..]);
+        let mut rset = RecordSet::default();
+        reader.read_record_set_limited(&mut rset, max_records);
+        assert_eq!(rset.len(), max_records);
+
+        let mut rset_iter = rset.into_iter();
+        let mut reader = Reader::new(&fasta_vec[..]);
+
+        for _ in 0..max_records {
+            let r0 = reader.next().unwrap().unwrap();
+            let rec = rset_iter.next().unwrap();
+            assert_eq!(rec.id(), r0.id());
+            assert_eq!(rec.desc(), r0.desc());
+            assert_eq!(rec.head(), r0.head());
+            assert_eq!(rec.seq(), r0.seq());
+        }
+    }
+}
+
 #[test]
 fn test_fasta_full_seq() {
     use std::borrow::Cow;

tests/fastq.rs

@@ -196,6 +196,35 @@ fn test_fastq_recset() {
     }
 }
 
+#[test]
+fn test_fastq_read_record_set_limited() {
+    for max_records in 3..10 {
+        let mut fastq_vec: Vec<u8> = Vec::with_capacity(FASTQ.len() * max_records);
+        for i in 0..max_records {
+            println!("i: {i} max_records: {max_records}");
+            fastq_vec.extend_from_slice(FASTQ);
+        }
+
+        let mut reader = Reader::new(&fastq_vec[..]);
+        let mut rset = RecordSet::default();
+        reader.read_record_set_limited(&mut rset, max_records);
+        assert_eq!(rset.len(), max_records);
+
+        let mut rset_iter = rset.into_iter();
+        let mut reader = Reader::new(&fastq_vec[..]);
+
+        for _ in 0..max_records {
+            let r0 = reader.next().unwrap().unwrap();
+            let rec = rset_iter.next().unwrap();
+            assert_eq!(rec.id(), r0.id());
+            assert_eq!(rec.desc(), r0.desc());
+            assert_eq!(rec.head(), r0.head());
+            assert_eq!(rec.seq(), r0.seq());
+            assert_eq!(rec.qual(), r0.qual());
+        }
+    }
+}
+
 #[test]
 fn test_fastq_parallel() {
     for cap in 3..400 {