nu_command/random/
dice.rs

1use nu_engine::command_prelude::*;
2use nu_protocol::ListStream;
3use rand::random_range;
4use std::num::NonZeroUsize;
5
6#[derive(Clone)]
7pub struct RandomDice;
8
9impl Command for RandomDice {
10    fn name(&self) -> &str {
11        "random dice"
12    }
13
14    fn signature(&self) -> Signature {
15        Signature::build("random dice")
16            .input_output_types(vec![(Type::Nothing, Type::list(Type::Int))])
17            .allow_variants_without_examples(true)
18            .named(
19                "dice",
20                SyntaxShape::Int,
21                "The amount of dice being rolled",
22                Some('d'),
23            )
24            .named(
25                "sides",
26                SyntaxShape::Int,
27                "The amount of sides a die has",
28                Some('s'),
29            )
30            .category(Category::Random)
31    }
32
33    fn description(&self) -> &str {
34        "Generate a random dice roll."
35    }
36
37    fn search_terms(&self) -> Vec<&str> {
38        vec!["generate", "die", "1-6"]
39    }
40
41    fn run(
42        &self,
43        engine_state: &EngineState,
44        stack: &mut Stack,
45        call: &Call,
46        _input: PipelineData,
47    ) -> Result<PipelineData, ShellError> {
48        dice(engine_state, stack, call)
49    }
50
51    fn examples(&self) -> Vec<Example> {
52        vec![
53            Example {
54                description: "Roll 1 dice with 6 sides each",
55                example: "random dice",
56                result: None,
57            },
58            Example {
59                description: "Roll 10 dice with 12 sides each",
60                example: "random dice --dice 10 --sides 12",
61                result: None,
62            },
63        ]
64    }
65}
66
67fn dice(
68    engine_state: &EngineState,
69    stack: &mut Stack,
70    call: &Call,
71) -> Result<PipelineData, ShellError> {
72    let span = call.head;
73
74    let sides: NonZeroUsize = call
75        .get_flag(engine_state, stack, "sides")?
76        .unwrap_or_else(|| NonZeroUsize::new(6).expect("default sides must be non-zero"));
77
78    let dice: NonZeroUsize = call
79        .get_flag(engine_state, stack, "dice")?
80        .unwrap_or_else(|| NonZeroUsize::new(1).expect("default dice count must be non-zero"));
81
82    let sides = sides.get();
83    let dice = dice.get();
84
85    let iter = (0..dice).map(move |_| Value::int(random_range(1..sides + 1) as i64, span));
86
87    Ok(ListStream::new(iter, span, engine_state.signals().clone()).into())
88}
89
90#[cfg(test)]
91mod test {
92    use super::*;
93
94    #[test]
95    fn test_examples() {
96        use crate::test_examples;
97
98        test_examples(RandomDice {})
99    }
100}