use std::io::{Write, Result};
use std::fs::File;
use std::path::Path;
use super::bgzip;
use super::{Header, Record, RecordWriter};
pub struct BamWriterBuilder {
write_header: bool,
level: u8,
additional_threads: u16,
}
impl BamWriterBuilder {
pub fn new() -> Self {
Self {
write_header: true,
level: 6,
additional_threads: 0,
}
}
pub fn write_header(&mut self, write: bool) -> &mut Self {
self.write_header = write;
self
}
pub fn compression_level(&mut self, level: u8) -> &mut Self {
assert!(level <= 9, "Maximal compression level is 9");
self.level = level;
self
}
pub fn additional_threads(&mut self, additional_threads: u16) -> &mut Self {
self.additional_threads = additional_threads;
self
}
pub fn from_path<P: AsRef<Path>>(&mut self, path: P, header: Header)
-> Result<BamWriter<File>> {
let stream = File::create(path)?;
self.from_stream(stream, header)
}
pub fn from_stream<W: Write>(&mut self, stream: W, header: Header) -> Result<BamWriter<W>> {
let mut writer = bgzip::Writer::build()
.additional_threads(self.additional_threads)
.compression_level(self.level)
.from_stream(stream);
if self.write_header {
header.write_bam(&mut writer)?;
}
writer.flush_contents()?;
Ok(BamWriter { writer, header })
}
}
pub struct BamWriter<W: Write> {
writer: bgzip::Writer<W>,
header: Header,
}
impl BamWriter<File> {
pub fn build() -> BamWriterBuilder {
BamWriterBuilder::new()
}
pub fn from_path<P: AsRef<Path>>(path: P, header: Header) -> Result<Self> {
BamWriterBuilder::new().from_path(path, header)
}
}
impl<W: Write> BamWriter<W> {
pub fn from_stream(stream: W, header: Header) -> Result<Self> {
BamWriterBuilder::new().from_stream(stream, header)
}
pub fn header(&self) -> &Header {
&self.header
}
pub fn take_stream(self) -> W {
self.writer.take_stream()
}
pub fn pause(&mut self) {
self.writer.pause();
}
}
impl<W: Write> RecordWriter for BamWriter<W> {
fn write(&mut self, record: &Record) -> Result<()> {
record.write_bam(&mut self.writer)?;
self.writer.end_context();
Ok(())
}
fn finish(&mut self) -> Result<()> {
self.writer.finish()
}
fn flush(&mut self) -> Result<()> {
self.writer.flush()
}
}