htsvcf_core/
writer.rs

1use crate::Header;
2use rust_htslib::bcf;
3use rust_htslib::errors::Result;
4use std::path::Path;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum OutputFormat {
8    Vcf,
9    Bcf,
10}
11
12#[derive(Debug, Clone, Default)]
13pub struct WriterOptions {
14    pub format: Option<OutputFormat>,
15    pub uncompressed: bool,
16    pub threads: Option<usize>,
17}
18
19#[derive(Debug)]
20pub struct Writer {
21    inner: bcf::Writer,
22}
23
24impl Writer {
25    pub fn write_record(&mut self, record: &mut bcf::Record) -> Result<()> {
26        self.inner.subset(record);
27        self.inner.write(record)
28    }
29}
30
31fn infer_format(path: &str) -> OutputFormat {
32    if path == "-" {
33        return OutputFormat::Vcf;
34    }
35    if path.ends_with(".bcf") {
36        OutputFormat::Bcf
37    } else {
38        OutputFormat::Vcf
39    }
40}
41
42fn to_rust_htslib_format(format: OutputFormat) -> bcf::Format {
43    match format {
44        OutputFormat::Vcf => bcf::Format::Vcf,
45        OutputFormat::Bcf => bcf::Format::Bcf,
46    }
47}
48
49fn clone_as_rust_htslib_header(header: &Header) -> bcf::header::Header {
50    let dup = unsafe { rust_htslib::htslib::bcf_hdr_dup(header.inner_ptr()) };
51    // SAFETY: dup is a fresh header owned by this header wrapper and will be freed
52    // by rust-htslib's Header Drop impl.
53    bcf::header::Header {
54        inner: dup,
55        subset: None,
56    }
57}
58
59pub fn open_writer(path: &str, header: &Header, opts: WriterOptions) -> Result<Writer> {
60    let format = opts.format.unwrap_or_else(|| infer_format(path));
61    let rust_header = clone_as_rust_htslib_header(header);
62
63    let inner = if path == "-" {
64        bcf::Writer::from_stdout(
65            &rust_header,
66            opts.uncompressed,
67            to_rust_htslib_format(format),
68        )?
69    } else {
70        bcf::Writer::from_path(
71            Path::new(path),
72            &rust_header,
73            opts.uncompressed,
74            to_rust_htslib_format(format),
75        )?
76    };
77
78    let mut out = Writer { inner };
79    if let Some(threads) = opts.threads {
80        if threads > 0 {
81            out.inner.set_threads(threads)?;
82        }
83    }
84
85    Ok(out)
86}