use std::path::PathBuf;
use anyhow::{Result, bail};
use clap::Args;
use libsail::collection::{Indexable, Iterable};
use rand::SeedableRng;
use rand::rngs::StdRng;
use crate::cli::FormatArg;
use crate::input::{Inputs, dispatch};
use crate::output::{emit, writer};
#[derive(Args)]
pub struct SampleArgs {
#[arg(default_value = "-")]
pub input: PathBuf,
#[arg(short = 'n', long, conflicts_with = "fraction")]
pub count: Option<usize>,
#[arg(short = 'p', long)]
pub fraction: Option<f64>,
#[arg(long)]
pub seed: Option<u64>,
#[arg(long, value_enum)]
pub format: Option<FormatArg>,
#[arg(short, long)]
pub output: Option<PathBuf>,
}
impl SampleArgs {
pub fn run(self) -> Result<()> {
let inputs = Inputs::read(std::slice::from_ref(&self.input), self.format)?;
let entry = &inputs.entries()[0];
let out = writer(self.output.as_deref())?;
let (count, fraction, seed) = (self.count, self.fraction, self.seed);
dispatch!(inputs.format(), entry, |collection, _size, _name, write| {
let n = draw_size(count, fraction, collection.len())?;
let mut rng = match seed {
Some(seed) => StdRng::seed_from_u64(seed),
None => StdRng::from_rng(&mut rand::rng()),
};
emit(collection.sample_with(&mut rng, n).iter(), write, out)
})
}
}
fn draw_size(count: Option<usize>, fraction: Option<f64>, len: usize) -> Result<usize> {
match (count, fraction) {
(Some(n), _) => {
if n > len {
bail!("{n} records were asked for and the input holds {len}")
}
Ok(n)
}
(None, Some(p)) => {
if !(0.0..=1.0).contains(&p) {
bail!("a fraction runs from 0 to 1, and {p} does not")
}
Ok((len as f64 * p).round() as usize)
}
(None, None) => bail!("sample needs -n or -p"),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn drawing_more_records_than_there_are_is_an_error_and_not_a_clamp() {
assert!(draw_size(Some(6), None, 5).is_err());
assert_eq!(draw_size(Some(5), None, 5).unwrap(), 5);
}
#[test]
fn a_fraction_rounds_rather_than_truncating() {
assert_eq!(draw_size(None, Some(0.5), 5).unwrap(), 3);
assert_eq!(draw_size(None, Some(1.0), 5).unwrap(), 5);
assert_eq!(draw_size(None, Some(0.0), 5).unwrap(), 0);
}
#[test]
fn a_fraction_outside_zero_to_one_is_refused() {
assert!(draw_size(None, Some(1.5), 5).is_err());
assert!(draw_size(None, Some(-0.5), 5).is_err());
}
#[test]
fn neither_n_nor_p_is_an_error_rather_than_an_empty_sample() {
assert!(draw_size(None, None, 5).is_err());
}
}