vcfkit 0.1.1

vcfkit: a simple program for vcf/bcf file manipulation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69

use colored::*;
use rust_htslib::bcf::{header::Header, Format, Read, Reader, Writer};
use rust_htslib::errors::Error;

pub fn select_sites(
    vcf: &Option<&str>,
    keep: &str,
    gzip: bool,
    out: &Option<&str>,
) -> Result<(), Error> {
    let mut bcf = if let Some(file) = vcf {
        Reader::from_path(file)?
    } else {
        Reader::from_stdin()?
    };

    let header1 = bcf.header();
    let header2 = Header::from_template(&header1);

    let mut bcf_writer = if let Some(file) = out {
        Writer::from_path(file, &header2, !gzip, Format::Vcf)?
    } else {
        Writer::from_stdout(&header2, !gzip, Format::Vcf)?
    };

    let mut flag_snp = 0;
    if keep == "SNP" {
        for rec in bcf.records() {
            if let Ok(ref info) = rec {
                for allele in info.alleles() {
                    if allele.len() == 1 {
                        continue;
                    } else {
                        flag_snp += 1;
                    }
                }
            }
            if flag_snp == 0 {
                bcf_writer.write(&rec?)?;
            }
            flag_snp = 0;
        }
    } else if keep == "INDEL" {
        for rec in bcf.records() {
            if let Ok(ref info) = rec {
                for allele in info.alleles() {
                    if allele.len() == 1 {
                        continue;
                    } else {
                        flag_snp += 1;
                    }
                }
            }
            if flag_snp != 0 {
                bcf_writer.write(&rec?)?;
            }
            flag_snp = 0;
        }
    } else {
        eprintln!(
            "{}",
            "[error]: invalid  arg '--keep', must be INDEL or SNP !".red()
        );
        std::process::exit(1);
    }

    Ok(())
}