forked from extern/nushell
65 lines
1.7 KiB
Rust
65 lines
1.7 KiB
Rust
use nu_errors::ShellError;
|
|
use nu_plugin::Plugin;
|
|
use nu_protocol::{
|
|
CallInfo, ReturnSuccess, ReturnValue, Signature, SyntaxShape, UntaggedValue, Value,
|
|
};
|
|
|
|
use rand::seq::SliceRandom;
|
|
use rand::thread_rng;
|
|
|
|
#[derive(Default)]
|
|
pub struct Shuffle {
|
|
values: Vec<ReturnValue>,
|
|
limit: Option<u64>,
|
|
}
|
|
|
|
impl Shuffle {
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
pub fn setup(&mut self, call_info: CallInfo) -> ReturnValue {
|
|
self.limit = if let Some(value) = call_info.args.get("num") {
|
|
Some(value.as_u64()?)
|
|
} else {
|
|
None
|
|
};
|
|
ReturnSuccess::value(UntaggedValue::nothing().into_untagged_value())
|
|
}
|
|
}
|
|
|
|
impl Plugin for Shuffle {
|
|
fn config(&mut self) -> Result<Signature, ShellError> {
|
|
Ok(Signature::build("shuffle")
|
|
.desc("Shuffle input randomly")
|
|
.named(
|
|
"num",
|
|
SyntaxShape::Int,
|
|
"Limit output to `num` number of values",
|
|
Some('n'),
|
|
)
|
|
.filter())
|
|
}
|
|
|
|
fn begin_filter(&mut self, call_info: CallInfo) -> Result<Vec<ReturnValue>, ShellError> {
|
|
self.setup(call_info)?;
|
|
Ok(vec![])
|
|
}
|
|
|
|
fn filter(&mut self, input: Value) -> Result<Vec<ReturnValue>, ShellError> {
|
|
self.values.push(input.into());
|
|
Ok(vec![])
|
|
}
|
|
|
|
fn end_filter(&mut self) -> Result<Vec<ReturnValue>, ShellError> {
|
|
let mut rng = thread_rng();
|
|
if let Some(n) = self.limit {
|
|
let (shuffled, _) = self.values.partial_shuffle(&mut rng, n as usize);
|
|
Ok(shuffled.to_vec())
|
|
} else {
|
|
self.values.shuffle(&mut rng);
|
|
Ok(self.values.clone())
|
|
}
|
|
}
|
|
}
|