use std::{
fs::File,
path::{Path, PathBuf},
time::Duration,
};
use anyhow::{Context, bail};
use bstr::ByteSlice;
use itertools::Itertools;
use rust_htslib::{
bam,
bcf::{
self, HeaderRecord, Read as _,
header::{HeaderView, TagLength, TagType},
},
};
use tracing::{debug, instrument, warn};
use url::Url;
use crate::{
RunnableCommand,
common::{
aligner::RUNNING,
csv::TwitcherCSVWriter,
list_of_regions::{Regions, RegionsDefinition, Targets},
reference::ReferenceReader,
},
vcf::{
cli::{Command, OutputOptions},
pipeline::{
VCFPipeline,
reader::VCFReader,
writer::{OutputWriter, region_writer::RegionWriter},
},
},
};
pub mod cli;
pub mod pipeline;
pub mod strings;
impl RunnableCommand for Command {
#[instrument(name = "vcf", skip_all)]
async fn run(self) -> anyhow::Result<()> {
tokio::spawn(async {
let mut interval = tokio::time::interval(Duration::from_secs(1));
loop {
interval.tick().await;
let r = RUNNING.load(std::sync::atomic::Ordering::Relaxed);
debug!("Running tasks: {r}");
}
});
self.validate()?;
let regions = self.regions.read_regions().await?;
let input = tokio::task::block_in_place(move || open_input_file(&self.input, regions))?;
warn_on_unwritable_format_fields(input.header());
let mut header = bcf::Header::from_template(input.header());
augment_header(&mut header, self.no_version);
let reference = tokio::task::block_in_place(|| ReferenceReader::try_from(&self.reference))?;
let targets = self.targets.read_regions().await?.map(Targets::new);
let region_output = RegionWriter::from_parameter(self.ts_regions.as_deref()).await?;
let output = tokio::task::block_in_place(move || open_output_file(&self.output, &header))?;
let csv_output = if let Some(f) = self.csv_output {
Some(TwitcherCSVWriter::new(Box::new(File::create(f)?)))
} else {
None
};
let bam = open_phasing_bam_file(self.phasing_bam.as_deref(), &self.reference.file)?;
VCFPipeline {
input,
output,
reference,
targets,
region_output,
csv_output,
settings: self.pipeline_settings,
bam,
phasing: self.phasing,
}
.run()
.await?;
Ok(())
}
}
impl Command {
fn validate(&self) -> anyhow::Result<()> {
if self.output.output_file.as_ref().is_none_or(|o| o == "-")
&& self.ts_regions.as_ref().is_some_and(|o| o == "-")
{
bail!("Cannot output both vcf and regions output on stdout.");
}
self.pipeline_settings.cluster.cluster_strategy.validate()?;
Ok(())
}
}
fn augment_header(header: &mut bcf::Header, no_version: bool) {
#[allow(clippy::wildcard_imports)]
use strings::*;
header.push_record(
format!(
"##INFO=<ID={VCF_TS_CIGARETS_KEY},Number=1,Type=String,Description=\"{VCF_TS_CIGARETS_DESC}\">"
)
.as_bytes(),
);
header.push_record(
format!(
"##INFO=<ID={VCF_TS_COST_KEY},Number=1,Type=Integer,Description=\"{VCF_TS_COST_DESC}\">"
)
.as_bytes(),
);
header.push_record(
format!(
"##INFO=<ID={VCF_TS_INNER_LEN_KEY},Number=.,Type=Integer,Description=\"{VCF_TS_INNER_LEN_DESC}\">"
)
.as_bytes(),
);
header.push_record(
format!(
"##INFO=<ID={VCF_TS_JUMP_KEY},Number=.,Type=Integer,Description=\"{VCF_TS_JUMP_DESC}\">"
)
.as_bytes(),
);
header.push_record(
format!(
"##INFO=<ID={VCF_TS_ID_KEY},Number=1,Type=String,Description=\"{VCF_TS_ID_DESC}\">"
)
.as_bytes(),
);
header.push_record(
format!(
"##INFO=<ID={VCF_CLUSTER_GRP_KEY},Number=1,Type=String,Description=\"{VCF_CLUSTER_GRP_DESC}\">"
)
.as_bytes(),
);
header.push_record(
format!(
"##INFO=<ID={VCF_TWITCHER_PHASE_KEY},Number=0,Type=Flag,Description=\"{VCF_TWITCHER_PHASE_DESC}\">"
)
.as_bytes(),
);
header.push_record(VCF_FORMAT_GT_LINE);
header.push_record(VCF_FORMAT_PS_LINE);
if !no_version {
header.push_record(
format!("##{VCF_HEADER_TWITCHER_VERSION_KEY}={TWITCHER_VERSION}").as_bytes(),
);
header.push_record(
format!("##{VCF_HEADER_TWITCHER_CMD}={}", std::env::args().join(" ")).as_bytes(),
);
}
}
fn warn_on_unwritable_format_fields(header: &HeaderView) {
for (tag, expected) in [
(&b"GT"[..], (TagType::String, TagLength::Fixed(1))),
(&b"PS"[..], (TagType::Integer, TagLength::Fixed(1))),
] {
if let Ok(found) = header.format_type(tag)
&& found != expected
{
warn!(
"Input header declares FORMAT/{} as {found:?}, but this program writes {expected:?}. \
The field will not be written.",
tag.as_bstr()
);
}
}
}
fn open_input_file(input: &str, regions: Option<Regions>) -> anyhow::Result<VCFReader> {
let reader = if let Some(regions) = regions {
let inner = match input {
"-" => bail!("Cannot create indexed reader from stdin"),
path_or_url => {
if let Ok(url) = Url::parse(path_or_url) {
bcf::IndexedReader::from_url(&url)
} else {
bcf::IndexedReader::from_path(path_or_url)
}
}
}?;
let regions =
regions.into_linear(|chr| inner.header().name2rid(chr.as_ref()).unwrap_or(0) as usize);
VCFReader::with_regions(inner, regions)
} else {
let inner = match input {
"-" => bcf::Reader::from_stdin(),
path_or_url => {
if let Ok(url) = Url::parse(path_or_url) {
bcf::Reader::from_url(&url)
} else {
bcf::Reader::from_path(path_or_url)
}
}
}?;
VCFReader::entire_file(inner)
};
let sample_count = reader.header().sample_count();
if sample_count > 1 {
let first_sample = reader
.header()
.samples()
.first()
.context("could not get sample name")?
.as_bstr();
warn!(
"The input file has {sample_count} samples, but twitcher only supports one sample per file. All but the first sample ({first_sample}) will be ignored."
);
}
if sample_count == 0 {
bail!(
"There needs to be a sample present in the input VCF file. Sites-only VCFs are currently not supported."
);
}
validate_version(get_vcf_version(reader.header())?)?;
Ok(reader)
}
fn validate_version((major, minor): (u8, u8)) -> anyhow::Result<()> {
if major != 4 || (!(2..5).contains(&minor)) {
bail!("VCF version unsupported. Twitcher only handles v4.2 <= version <= v4.5");
}
Ok(())
}
fn get_vcf_version(header: &HeaderView) -> anyhow::Result<(u8, u8)> {
let records = header.header_records();
let first = records
.first()
.with_context(|| "Header does not have at least one record.")?;
if let HeaderRecord::Generic { key, value } = first
&& key == "fileformat"
&& value.starts_with("VCFv")
{
let version = value
.get("VCFv".len()..)
.context("version string not long enough")?;
let Some((major, minor)) = version.split_once('.') else {
anyhow::bail!("no dot separating major and minor version");
};
return Ok((major.parse()?, minor.parse()?));
}
bail!("The first header record of the VCF input file must be the VCF version.");
}
fn open_phasing_bam_file<P: AsRef<Path>>(
arg: Option<&str>,
reference_path: P,
) -> anyhow::Result<Option<bam::IndexedReader>> {
let Some(path) = arg else {
return Ok(None);
};
let mut reader = if let Ok(url) = Url::parse(path) {
bam::IndexedReader::from_url(&url)?
} else {
bam::IndexedReader::from_path(path)?
};
reader.set_reference(reference_path)?;
Ok(Some(reader))
}
fn open_output_file(output: &OutputOptions, header: &bcf::Header) -> anyhow::Result<OutputWriter> {
let inner = match output.output_file.as_deref() {
None | Some("-") => bcf::Writer::from_stdout(header, true, bcf::Format::Vcf),
Some(path_or_url) => {
let url = Url::parse(path_or_url);
let path = PathBuf::from(match &url {
Ok(url) => url.path(),
_ => path_or_url,
});
let compressed = matches!(
path.extension().and_then(|e| e.to_str()),
Some("gz" | "bgz")
);
let format = if path
.file_name()
.is_none_or(|n| n.to_str().is_some_and(|s| s.contains(".vcf")))
{
bcf::Format::Vcf
} else {
bcf::Format::Bcf
};
match url {
Ok(url) => bcf::Writer::from_url(&url, header, !compressed, format),
Err(_) => bcf::Writer::from_path(path, header, !compressed, format),
}
}
}?;
let writer = if let Some(plus_minus) = output.output_filter.realigned_and_context {
OutputWriter::new_buffered(inner, i64::try_from(plus_minus)?)
} else {
OutputWriter::new_native(inner, output.output_filter.only_realigned)
};
Ok(writer)
}
#[cfg(test)]
mod tests {
use std::io::Read as _;
use rust_htslib::{bcf::Header, bgzf};
use crate::common::reference::CliReferenceArg;
use super::*;
#[test]
fn validate_rejects_both_outputs_on_stdout() {
let cmd = Command {
output: None::<&str>.into(), ts_regions: Some("-".into()),
..Default::default()
};
assert!(cmd.validate().is_err());
}
#[test]
fn validate_allows_vcf_on_file_regions_on_stdout() {
let cmd = Command {
output: Some("/tmp/out.vcf").into(),
ts_regions: Some("-".into()),
..Default::default()
};
assert!(cmd.validate().is_ok());
}
#[tokio::test(flavor = "multi_thread")]
async fn test_pipeline() {
let dir = tempfile::tempdir().unwrap();
let outpath = format!("{}/out.vcf", dir.path().to_str().unwrap());
let cmd = Command {
input: "./test_files/test.vcf".to_string(),
reference: CliReferenceArg::from("./test_files/test.fa"),
output: Some(outpath.clone()).into(),
..Default::default()
};
cmd.run().await.unwrap();
let output = std::fs::read_to_string(outpath).unwrap();
let out_records = output.lines().filter(|l| !l.starts_with('#'));
assert_eq!(out_records.count(), 2);
}
#[test]
fn test_auto_detect_output_format_vcf() {
let reader = open_input_file("./test_files/test.vcf", None).unwrap();
let header = Header::from_template(reader.header());
let dir = tempfile::tempdir().unwrap();
let outpath_vcf = format!("{}/out.vcf", dir.path().to_str().unwrap());
let writer = open_output_file(&Some(&outpath_vcf).into(), &header).unwrap();
drop(writer);
let written_file = std::fs::read_to_string(outpath_vcf).unwrap();
assert!(written_file.starts_with("##fileformat=VCFv4"));
}
#[test]
fn test_auto_detect_output_format_vcf_gz() {
let reader = open_input_file("./test_files/test.vcf", None).unwrap();
let header = Header::from_template(reader.header());
let dir = tempfile::tempdir().unwrap();
let outpath_vcf_gz = format!("{}/out.vcf.gz", dir.path().to_str().unwrap());
let writer = open_output_file(&Some(&outpath_vcf_gz).into(), &header).unwrap();
drop(writer);
let mut written_file = Vec::new();
bgzf::Reader::from_path(outpath_vcf_gz)
.unwrap()
.read_to_end(&mut written_file)
.unwrap();
assert!(written_file.starts_with(b"##fileformat=VCFv4"));
}
#[test]
fn test_auto_detect_output_format_bcf() {
let reader = open_input_file("./test_files/test.vcf", None).unwrap();
let header = Header::from_template(reader.header());
let dir = tempfile::tempdir().unwrap();
let outpath_bcf = format!("{}/out.bcf", dir.path().to_str().unwrap());
let writer = open_output_file(&Some(&outpath_bcf).into(), &header).unwrap();
drop(writer);
let mut written_file = Vec::new();
bgzf::Reader::from_path(outpath_bcf)
.unwrap()
.read_to_end(&mut written_file)
.unwrap();
assert!(written_file.starts_with(b"BCF"));
}
#[test]
fn test_auto_detect_output_format_bcf_gz() {
let reader = open_input_file("./test_files/test.vcf", None).unwrap();
let header = Header::from_template(reader.header());
let dir = tempfile::tempdir().unwrap();
let outpath_bcf_gz = format!("{}/out.bcf.gz", dir.path().to_str().unwrap());
let writer = open_output_file(&Some(&outpath_bcf_gz).into(), &header).unwrap();
drop(writer);
let mut written_file = Vec::new();
bgzf::Reader::from_path(outpath_bcf_gz)
.unwrap()
.read_to_end(&mut written_file)
.unwrap();
assert!(written_file.starts_with(b"BCF"));
}
}