fix(random dice): gracefully handle --sides 0 using NonZeroUsize (#16001)

This commit is contained in:
Kumar Ujjawal
2025-06-23 18:17:50 +05:30
committed by GitHub
parent 680a2fa2aa
commit ae0cf8780d
2 changed files with 51 additions and 2 deletions

View File

@ -1,6 +1,7 @@
use nu_engine::command_prelude::*;
use nu_protocol::ListStream;
use rand::random_range;
use std::num::NonZeroUsize;
#[derive(Clone)]
pub struct RandomDice;
@ -70,8 +71,16 @@ fn dice(
) -> Result<PipelineData, ShellError> {
let span = call.head;
let dice: usize = call.get_flag(engine_state, stack, "dice")?.unwrap_or(1);
let sides: usize = call.get_flag(engine_state, stack, "sides")?.unwrap_or(6);
let sides: NonZeroUsize = call
.get_flag(engine_state, stack, "sides")?
.unwrap_or_else(|| NonZeroUsize::new(6).expect("default sides must be non-zero"));
let dice: NonZeroUsize = call
.get_flag(engine_state, stack, "dice")?
.unwrap_or_else(|| NonZeroUsize::new(1).expect("default dice count must be non-zero"));
let sides = sides.get();
let dice = dice.get();
let iter = (0..dice).map(move |_| Value::int(random_range(1..sides + 1) as i64, span));

View File

@ -10,6 +10,10 @@ use std::{
cmp::Ordering,
collections::{HashMap, VecDeque},
fmt,
num::{
NonZeroI8, NonZeroI16, NonZeroI32, NonZeroI64, NonZeroIsize, NonZeroU16, NonZeroU32,
NonZeroU64, NonZeroUsize,
},
path::PathBuf,
str::FromStr,
};
@ -270,6 +274,42 @@ impl FromValue for i64 {
}
}
//
// We can not use impl<T: FromValue> FromValue for NonZero<T> as NonZero requires an unstable trait
// As a result, we use this macro to implement FromValue for each NonZero type.
//
macro_rules! impl_from_value_for_nonzero {
($nonzero:ty, $base:ty) => {
impl FromValue for $nonzero {
fn from_value(v: Value) -> Result<Self, ShellError> {
let span = v.span();
let val = <$base>::from_value(v)?;
<$nonzero>::new(val).ok_or_else(|| ShellError::IncorrectValue {
msg: "use a value other than 0".into(),
val_span: span,
call_span: span,
})
}
fn expected_type() -> Type {
Type::Int
}
}
};
}
impl_from_value_for_nonzero!(NonZeroU16, u16);
impl_from_value_for_nonzero!(NonZeroU32, u32);
impl_from_value_for_nonzero!(NonZeroU64, u64);
impl_from_value_for_nonzero!(NonZeroUsize, usize);
impl_from_value_for_nonzero!(NonZeroI8, i8);
impl_from_value_for_nonzero!(NonZeroI16, i16);
impl_from_value_for_nonzero!(NonZeroI32, i32);
impl_from_value_for_nonzero!(NonZeroI64, i64);
impl_from_value_for_nonzero!(NonZeroIsize, isize);
macro_rules! impl_from_value_for_int {
($type:ty) => {
impl FromValue for $type {