瀏覽代碼

save progress

Thomas 1 年之前
父節點
當前提交
5e0aa1987a
共有 2 個文件被更改,包括 42 次插入4 次删除
  1. 37 3
      src/in_out/mod.rs
  2. 5 1
      src/variants.rs

+ 37 - 3
src/in_out/mod.rs

@@ -1,10 +1,16 @@
-use std::{fs::File, io::BufReader};
-use anyhow::{ Ok, Result};
+use anyhow::{Ok, Result};
 use bgzip::BGZFReader;
+use indicatif::MultiProgress;
+use std::{
+    fs::{File, Metadata},
+    io::BufReader,
+};
 
+use crate::utils::new_pg_speed;
+
+pub mod dict_reader;
 pub mod vcf_reader;
 pub mod vcf_writer;
-pub mod dict_reader;
 
 pub fn get_reader(path: &str) -> Result<Box<dyn std::io::Read>> {
     let file_type = *path.split(".").collect::<Vec<&str>>().last().unwrap();
@@ -27,4 +33,32 @@ pub fn get_reader(path: &str) -> Result<Box<dyn std::io::Read>> {
         }
     }
 }
+pub fn get_reader_progress(
+    path: &str,
+    mp: &MultiProgress,
+) -> Result<Box<dyn std::io::Read>> {
+    let file_type = *path.split(".").collect::<Vec<&str>>().last().unwrap();
+
+    assert!(file_type == "gz" || file_type == "vcf");
+    let file = File::open(path)?;
+    let metadata = file.metadata()?;
+    let pg = mp.add(new_pg_speed(metadata.len() as u64));
+    pg.set_message(format!("Reading {path}"));
 
+    let raw_reader: Box<dyn std::io::Read> = Box::new(file);
+    let raw_reader = pg.wrap_read(raw_reader);
+
+    match file_type {
+        "gz" => {
+            let reader = Box::new(BGZFReader::new(raw_reader)?);
+            Ok(Box::new(BufReader::new(reader)))
+        }
+        "vcf" => {
+            // let reader = Box::new(BzDecoder::new(raw_reader));
+            Ok(Box::new(BufReader::new(raw_reader)))
+        }
+        t => {
+            panic!("unknown file type: {}", t)
+        }
+    }
+}

+ 5 - 1
src/variants.rs

@@ -857,7 +857,11 @@ impl Variants {
     }
 
     pub fn new_from_bytes(name: &str, path: &str, mp: MultiProgress) -> Result<Self> {
-        let r = in_out::get_reader(path)?;
+        info!("Loading variants from: {path}");
+        let r = in_out::get_reader_progress(path, &mp)?;
+        
+
+        
         let data: Vec<Variant> = pot::from_reader(r)?;
         Ok(Self {
             name: name.to_string(),