use ahda::EncodeOpts;
use std::fs::File;
use std::io::BufRead;
use std::io::BufReader;
use std::io::Read;
use std::io::Write;
use std::io::IsTerminal;
use std::path::PathBuf;
use std::os::unix::fs::MetadataExt;
use std::os::unix::fs::FileTypeExt;
use clap::Parser;
mod cli;
fn init_log(log_max_level: usize) {
stderrlog::new()
.module(module_path!())
.quiet(false)
.verbosity(log_max_level)
.timestamp(stderrlog::Timestamp::Off)
.init()
.unwrap();
}
type E = Box<dyn std::error::Error>;
struct FastxNameReader {
reader: Box<dyn needletail::FastxReader>,
}
impl FastxNameReader {
pub fn new(
file: &PathBuf,
) -> Result<Self, E> {
let reader = needletail::parse_fastx_file(file)?;
Ok(Self{ reader })
}
}
impl Iterator for FastxNameReader {
type Item = Vec<u8>;
fn next(
&mut self,
) -> Option<Vec<u8>> {
let record = self.reader.next()?;
let query_info = record.unwrap();
let end = query_info.id().iter().position(|x| x == &b' ');
Some(query_info.id()[0..end.unwrap_or(query_info.id().len())].to_vec())
}
}
fn stdin_is_piped_data() -> bool {
if std::io::stdin().is_terminal() {
false
} else {
std::fs::metadata("/dev/stdin")
.map(|meta| {
let file_type = meta.file_type();
file_type.is_fifo() || (file_type.is_file() && meta.size() > 0)
})
.unwrap_or(false)
}
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
let cli = cli::Cli::parse();
match &cli.command {
Some(cli::Commands::Encode {
input_file,
input_format,
query_file,
target_list,
sample_name,
stdout,
force,
keep,
rename,
verbose,
}) => {
init_log(if *verbose { 2 } else { 1 });
let mut targets = None;
if let Some(target_list) = target_list {
match File::open(target_list) {
Ok(f) => {
let reader = BufReader::new(f);
targets = Some(reader.split(b'\n').map(|x| x.unwrap()).collect::<Vec<Vec<u8>>>());
},
Err(e) => {
eprintln!("ahda: can't open input file `{}`: {}", target_list.to_string_lossy(), e);
return Err(Box::new(e))
},
}
}
let queries: Option<FastxNameReader> = if let Some(query_file) = query_file {
match FastxNameReader::new(query_file) {
Ok(reader) => Some(reader),
Err(e) => {
eprintln!("ahda: can't open input file `{}`: {}", query_file.to_string_lossy(), e);
return Err(e)
},
}
} else {
None
};
let mut inputs: Vec<Box<dyn Read>> = Vec::new();
let mut outputs: Vec<Box<dyn Write>> = Vec::new();
let mut force_stdout: bool = false;
if let Some(input_file) = input_file {
match File::open(input_file) {
Ok(conn_in) => inputs.push(Box::new(conn_in)),
Err(e) => {
eprintln!("ahda: can't open input file `{}`: {}", input_file.to_string_lossy(), e);
return Err(Box::new(e))
},
}
let out_path = PathBuf::from(input_file.to_string_lossy().to_string() + ".ahda");
if !*stdout {
match if *force { File::create(out_path.clone()) } else { File::create_new(out_path.clone()) } {
Ok(conn_out) => {
outputs.push(Box::new(conn_out));
},
Err(e) => {
eprintln!("ahda: can't create output file `{}`: {}", out_path.to_string_lossy(), e);
return Err(Box::new(e))
},
}
}
} else {
inputs.push(Box::new(std::io::stdin()));
if !*force && std::io::stdout().is_terminal() {
eprintln!("ahda: refusing to write binary data to terminal, use `--force` to override");
return Ok(());
} else {
force_stdout = true;
}
}
if *stdout || force_stdout {
outputs.push(Box::new(std::io::stdout()));
}
let mut opts = EncodeOpts::default();
opts.format = input_format.clone();
opts.encode_query_names = query_file.is_some() && !*rename;
opts.rename_queries = *rename;
opts.accession = if let Some(name) = sample_name {
name.as_bytes().to_vec()
} else {
if let Some(file) = query_file {
file.to_string_lossy().as_bytes().to_vec()
} else {
eprintln!("ahda: use `--name` to supply the sample name");
return Ok(())
}
};
let conn_in = &mut inputs[0];
let conn_out = &mut outputs[0];
#[allow(clippy::manual_map)]
let t_it = if let Some(t) = targets { Some(&mut t.into_iter()) } else { None };
let ret = if let Some(mut q_it) = queries {
ahda::encode_from_read_to_write(t_it, Some(&mut q_it), conn_in, conn_out, opts)
} else {
ahda::encode_from_read_to_write(t_it, None::<&mut std::iter::Empty<Vec<u8>>>, conn_in, conn_out, opts)
};
if let Err(e) = ret {
eprintln!("ahda: can't encode input file `{}`: {}", input_file.as_ref().unwrap().to_string_lossy(), e);
return Err(e)
}
if !*keep && !*stdout && input_file.is_some() {
match std::fs::remove_file(input_file.as_ref().unwrap()) {
Ok(()) => (),
Err(e) => {
eprintln!("ahda: can't remove input file `{}`: {}", input_file.as_ref().unwrap().to_string_lossy(), e);
return Err(Box::new(e))
}
}
}
Ok(())
},
Some(cli::Commands::Decode {
input_file,
format,
stdout,
force,
keep,
verbose,
}) => {
init_log(if *verbose { 2 } else { 1 });
let mut conn_in: Vec<Box<dyn Read>> = Vec::new();
let mut conn_out: Vec<Box<dyn Write>> = Vec::new();
if let Some(file) = input_file {
if *stdout {
conn_out.push(Box::new(std::io::stdout()));
} else {
let out_name = file.file_stem().unwrap().to_string_lossy();
let out_path = PathBuf::from(out_name.to_string());
if !*stdout {
match if *force { File::create(out_path.clone()) } else { File::create_new(out_path.clone()) } {
Ok(out) => {
conn_out.push(Box::new(out));
},
Err(e) => {
eprintln!("ahda: can't create output file `{}`: {}", out_path.to_string_lossy(), e);
return Err(Box::new(e))
},
}
}
}
match File::open(file) {
Ok(conn) => conn_in.push(Box::new(conn)),
Err(e) => {
eprintln!("ahda: can't open input file `{}`: {}", file.to_string_lossy(), e);
return Err(Box::new(e))
},
}
} else {
if !*force && std::io::stdin().is_terminal() {
eprintln!("ahda: standard input is a terminal, use `--force` to ignore");
return Ok(());
}
conn_in.push(Box::new(std::io::stdin()));
conn_out.push(Box::new(std::io::stdout()));
}
ahda::decode_from_read_to_write(format.clone().unwrap_or_default(), &mut conn_in[0], &mut conn_out[0]).unwrap();
if !*keep && !*stdout && input_file.is_some() {
match std::fs::remove_file(input_file.as_ref().unwrap()) {
Ok(()) => (),
Err(e) => {
eprintln!("ahda: can't remove input file `{}`: {}", input_file.as_ref().unwrap().to_string_lossy(), e);
return Err(Box::new(e))
}
}
}
Ok(())
},
Some(cli::Commands::Cat {
input_files,
output_file,
stdout,
force,
verbose,
}) => {
init_log(if *verbose { 2 } else { 1 });
assert!(!input_files.is_empty());
let mut conn_in: Vec<Box<dyn Read>> = Vec::new();
let has_piped_data = stdin_is_piped_data();
if has_piped_data {
conn_in.push(Box::new(std::io::stdin()));
}
for file in input_files {
match File::open(file) {
Ok(conn) => conn_in.push(Box::new(conn)),
Err(e) => {
eprintln!("ahda: can't open input file `{}`: {}", file.to_string_lossy(), e);
return Err(Box::new(e))
},
}
}
let mut conn_out: Vec<Box<dyn Write>> = Vec::new();
if let Some(file) = output_file {
if !*stdout {
match if *force { File::create(file.clone()) } else { File::create_new(file.clone()) } {
Ok(out) => {
conn_out.push(Box::new(out));
},
Err(e) => {
eprintln!("ahda: can't create output file `{}`: {}", file.to_string_lossy(), e);
return Err(Box::new(e))
},
}
}
}
if (*stdout || (output_file.is_none())) && *force {
conn_out.push(Box::new(std::io::stdout()));
} else if !*force && std::io::stdout().is_terminal() {
eprintln!("ahda: refusing to write binary data to terminal, use `--force` to override");
return Ok(());
}
let opts = EncodeOpts::default();
match ahda::concatenate_from_read_to_write(&mut conn_in, &mut conn_out[0], opts) {
Ok(_) => Ok(()),
Err(e) => {
let mut msg = input_files.iter().map(|x| x.to_string_lossy() + " ").collect::<String>();
msg.remove(msg.len() - 1);
eprintln!("ahda: can't concatenate input files `{}`: {}", msg, e);
Err(e)
}
}
},
Some(cli::Commands::Convert {
input_file,
query_file,
target_list,
output_file,
format,
sample_name,
stdout,
force,
keep,
verbose,
}) => {
init_log(if *verbose { 2 } else { 1 });
let mut targets = None;
if let Some(target_list) = target_list {
match File::open(target_list) {
Ok(f) => {
let reader = BufReader::new(f);
targets = Some(reader.split(b'\n').map(|x| x.unwrap()).collect::<Vec<Vec<u8>>>());
},
Err(e) => {
eprintln!("ahda: can't open input file `{}`: {}", target_list.to_string_lossy(), e);
return Err(Box::new(e))
},
}
}
let queries: Option<FastxNameReader> = if let Some(query_file) = query_file {
match FastxNameReader::new(query_file) {
Ok(reader) => Some(reader),
Err(e) => {
eprintln!("ahda: can't open input file `{}`: {}", query_file.to_string_lossy(), e);
return Err(e)
},
}
} else {
None
};
let mut inputs: Vec<Box<dyn Read>> = Vec::new();
let mut outputs: Vec<Box<dyn Write>> = Vec::new();
if let Some(input_file) = input_file {
match File::open(input_file) {
Ok(conn_in) => inputs.push(Box::new(conn_in)),
Err(e) => {
eprintln!("ahda: can't open input file `{}`: {}", input_file.to_string_lossy(), e);
return Err(Box::new(e))
},
}
let out_path = if let Some(out) = output_file { out.clone() } else { PathBuf::from(input_file.to_string_lossy().to_string() + ".tmp") };
if !*stdout {
match if *force { File::create(out_path.clone()) } else { File::create_new(out_path.clone()) } {
Ok(conn_out) => {
outputs.push(Box::new(conn_out));
},
Err(e) => {
eprintln!("ahda: can't create output file `{}`: {}", out_path.to_string_lossy(), e);
return Err(Box::new(e))
},
}
}
} else {
if std::io::stdin().is_terminal() {
eprintln!("ahda: standard input is a terminal, ignoring");
return Ok(());
}
inputs.push(Box::new(std::io::stdin()));
}
if *stdout || (input_file.is_none() && output_file.is_none()) {
outputs.push(Box::new(std::io::stdout()));
}
let conn_in = &mut inputs[0];
let conn_out = &mut outputs[0];
#[allow(clippy::manual_map)]
let t_it = if let Some(t) = targets { Some(&mut t.into_iter()) } else { None };
let ret = if let Some(mut q_it) = queries {
let sample = if let Some(name) = sample_name { name.as_bytes().to_vec() } else { query_file.as_ref().unwrap().to_string_lossy().as_bytes().to_vec() };
ahda::convert_from_read_to_write(t_it, Some(&mut q_it), &sample, format.as_ref().unwrap().clone(), conn_in, conn_out)
} else {
let sample = if let Some(name) = sample_name { name.as_bytes().to_vec() } else {
eprintln!("ahda: use `--name` to supply the sample name");
return Ok(())
};
ahda::convert_from_read_to_write(t_it, None::<&mut std::iter::Empty<Vec<u8>>>, &sample, format.as_ref().unwrap().clone(), conn_in, conn_out)
};
if let Err(e) = ret {
eprintln!("ahda: can't convert input file `{}`: {}", input_file.as_ref().unwrap().to_string_lossy(), e);
return Err(e)
}
if !*keep && !*stdout && input_file.is_some() && output_file.is_none() {
let tmp_path = PathBuf::from(input_file.as_ref().unwrap().to_string_lossy().to_string() + ".tmp");
match std::fs::rename(tmp_path.clone(), input_file.as_ref().unwrap()) {
Ok(()) => (),
Err(e) => {
eprintln!("ahda: can't rename output file `{}` to `{}`: {}", input_file.as_ref().unwrap().to_string_lossy(), tmp_path.to_string_lossy(), e);
return Err(Box::new(e))
}
}
}
Ok(())
},
Some(cli::Commands::Set {
input_files,
output_file,
operation,
stdout,
force,
rename,
strip_names,
sample_name,
verbose,
}) => {
init_log(if *verbose { 2 } else { 1 });
assert!(!input_files.is_empty());
let mut conn_in: Vec<Box<dyn Read>> = Vec::new();
let has_piped_data = stdin_is_piped_data();
if has_piped_data {
conn_in.push(Box::new(std::io::stdin()));
}
for file in input_files {
match File::open(file) {
Ok(conn) => conn_in.push(Box::new(conn)),
Err(e) => {
eprintln!("ahda: can't open input file `{}`: {}", file.to_string_lossy(), e);
return Err(Box::new(e))
},
}
}
let mut conn_out: Vec<Box<dyn Write>> = Vec::new();
if let Some(file) = output_file {
if !*stdout {
match if *force { File::create(file.clone()) } else { File::create_new(file.clone()) } {
Ok(out) => {
conn_out.push(Box::new(out));
},
Err(e) => {
eprintln!("ahda: can't create output file `{}`: {}", file.to_string_lossy(), e);
return Err(Box::new(e))
},
}
}
}
if *stdout || (output_file.is_none()) {
conn_out.push(Box::new(std::io::stdout()));
}
let mut opts = EncodeOpts::default();
opts.encode_query_names = !*strip_names;
opts.rename_queries = *rename;
if let Some(name) = sample_name {
opts.accession = name.as_bytes().to_vec();
}
opts.metadata_compression = ahda::compression::MetadataCompression::Flate2;
match ahda::merge_from_read_to_write(&mut conn_in, &mut conn_out[0], operation.as_ref().unwrap(), opts) {
Ok(_) => Ok(()),
Err(e) => {
eprintln!("ahda: could not merge input files: {}", e);
Err(e)
}
}
},
None => { eprintln!("ahda: Try 'ahda --help' for more information."); Ok(()) },
}
}