Skip to main content

nu_command/math/
cbrt.rs

1use crate::math::utils::run_with_elementwise;
2use nu_engine::command_prelude::*;
3
4#[derive(Clone)]
5pub struct MathCbrt;
6
7impl Command for MathCbrt {
8    fn name(&self) -> &str {
9        "math cbrt"
10    }
11
12    fn signature(&self) -> Signature {
13        Signature::build("math cbrt")
14            .input_output_types(vec![
15                (Type::Number, Type::Float),
16                (
17                    Type::List(Box::new(Type::Number)),
18                    Type::List(Box::new(Type::Float)),
19                ),
20                (Type::Range, Type::List(Box::new(Type::Number))),
21                (Type::record(), Type::record()),
22            ])
23            .rest(
24                "columns",
25                SyntaxShape::CellPath,
26                "The cell-paths/columns to operate on.",
27            )
28            .allow_variants_without_examples(true)
29            .category(Category::Math)
30    }
31
32    fn description(&self) -> &str {
33        "Returns the real-valued cube root of the input number."
34    }
35
36    fn search_terms(&self) -> Vec<&str> {
37        vec!["cube", "root"]
38    }
39
40    fn is_const(&self) -> bool {
41        true
42    }
43
44    fn run(
45        &self,
46        engine_state: &EngineState,
47        stack: &mut Stack,
48        call: &Call,
49        input: PipelineData,
50    ) -> Result<PipelineData, ShellError> {
51        let cell_paths: Vec<CellPath> = call.rest(engine_state, stack, 0)?;
52        let head = call.head;
53        run_with_elementwise(
54            input,
55            cell_paths,
56            head,
57            engine_state.signals(),
58            true,
59            move |value| operate(value, head),
60        )
61    }
62
63    fn run_const(
64        &self,
65        working_set: &StateWorkingSet,
66        call: &Call,
67        input: PipelineData,
68    ) -> Result<PipelineData, ShellError> {
69        let cell_paths: Vec<CellPath> = call.rest_const(working_set, 0)?;
70        let head = call.head;
71        run_with_elementwise(
72            input,
73            cell_paths,
74            head,
75            working_set.permanent().signals(),
76            true,
77            move |value| operate(value, head),
78        )
79    }
80
81    fn examples(&self) -> Vec<Example<'_>> {
82        vec![
83            Example {
84                description: "Compute the cube root of each number in a list.",
85                example: "[8 -27] | math cbrt",
86                result: Some(Value::list(
87                    vec![Value::test_float(2.0), Value::test_float(-3.0)],
88                    Span::test_data(),
89                )),
90            },
91            Example {
92                description: "Compute the cube root of list-valued columns in a record.",
93                example: "{alice: [8 27 64], bob: [125 216]} | math cbrt",
94                result: Some(Value::test_record(record! {
95                    "alice" => Value::list(
96                        vec![Value::test_float(2.0), Value::test_float(3.0), Value::test_float(4.0)],
97                        Span::test_data(),
98                    ),
99                    "bob" => Value::list(
100                        vec![Value::test_float(5.0), Value::test_float(6.0)],
101                        Span::test_data(),
102                    ),
103                })),
104            },
105            Example {
106                description: "Compute the cube root of a single column using a cell path.",
107                example: "{alice: [8 27 64], bob: [125 216]} | math cbrt alice",
108                result: Some(Value::test_record(record! {
109                    "alice" => Value::list(
110                        vec![Value::test_float(2.0), Value::test_float(3.0), Value::test_float(4.0)],
111                        Span::test_data(),
112                    ),
113                    "bob" => Value::list(
114                        vec![Value::test_int(125), Value::test_int(216)],
115                        Span::test_data(),
116                    ),
117                })),
118            },
119        ]
120    }
121}
122
123fn operate(value: Value, head: Span) -> Value {
124    let span = value.span();
125    match value {
126        Value::Int { val, .. } => Value::float((val as f64).cbrt(), span),
127        Value::Float { val, .. } => Value::float(val.cbrt(), span),
128        Value::Error { .. } => value,
129        other => Value::error(
130            ShellError::OnlySupportsThisInputType {
131                exp_input_type: crate::math::utils::NUMBER_INPUT_TYPES.into(),
132                wrong_type: other.get_type().to_string(),
133                dst_span: head,
134                src_span: other.span(),
135            },
136            head,
137        ),
138    }
139}
140
141#[cfg(test)]
142mod test {
143    use super::*;
144
145    #[test]
146    fn test_examples() -> nu_test_support::Result {
147        nu_test_support::test().examples(MathCbrt)
148    }
149}