nu_command/math/
sqrt.rs

1use crate::math::utils::run_with_elementwise;
2use nu_engine::command_prelude::*;
3
4#[derive(Clone)]
5pub struct MathSqrt;
6
7impl Command for MathSqrt {
8    fn name(&self) -> &str {
9        "math sqrt"
10    }
11
12    fn signature(&self) -> Signature {
13        Signature::build("math sqrt")
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 square root of the input number."
34    }
35
36    fn search_terms(&self) -> Vec<&str> {
37        vec!["square", "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 square root of each number in a list.",
85                example: "[9 16] | math sqrt",
86                result: Some(Value::list(
87                    vec![Value::test_float(3.0), Value::test_float(4.0)],
88                    Span::test_data(),
89                )),
90            },
91            Example {
92                description: "Apply square root to list-valued columns in a record.",
93                example: "{alice: [1 4 9], bob: [16 25 36]} | math sqrt",
94                result: Some(Value::test_record(record! {
95                    "alice" => Value::list(
96                        vec![Value::test_float(1.0), Value::test_float(2.0), Value::test_float(3.0)],
97                        Span::test_data(),
98                    ),
99                    "bob" => Value::list(
100                        vec![Value::test_float(4.0), Value::test_float(5.0), Value::test_float(6.0)],
101                        Span::test_data(),
102                    ),
103                })),
104            },
105            Example {
106                description: "Apply square root to a single column using a cell path.",
107                example: "{alice: [1 4 9], bob: [16 25 36]} | math sqrt alice",
108                result: Some(Value::test_record(record! {
109                    "alice" => Value::list(
110                        vec![Value::test_float(1.0), Value::test_float(2.0), Value::test_float(3.0)],
111                        Span::test_data(),
112                    ),
113                    "bob" => Value::list(
114                        vec![Value::test_int(16), Value::test_int(25), Value::test_int(36)],
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, .. } => {
127            let squared = (val as f64).sqrt();
128            if squared.is_nan() {
129                return error_negative_sqrt(head, span);
130            }
131            Value::float(squared, span)
132        }
133        Value::Float { val, .. } => {
134            let squared = val.sqrt();
135            if squared.is_nan() {
136                return error_negative_sqrt(head, span);
137            }
138            Value::float(squared, span)
139        }
140        Value::Error { .. } => value,
141        other => Value::error(
142            ShellError::OnlySupportsThisInputType {
143                exp_input_type: crate::math::utils::NUMBER_INPUT_TYPES.into(),
144                wrong_type: other.get_type().to_string(),
145                dst_span: head,
146                src_span: other.span(),
147            },
148            head,
149        ),
150    }
151}
152
153fn error_negative_sqrt(head: Span, span: Span) -> Value {
154    Value::error(
155        ShellError::UnsupportedInput {
156            msg: String::from("Can't square root a negative number"),
157            input: "value originates from here".into(),
158            msg_span: head,
159            input_span: span,
160        },
161        span,
162    )
163}
164
165#[cfg(test)]
166mod test {
167    use super::*;
168
169    #[test]
170    fn test_examples() -> nu_test_support::Result {
171        nu_test_support::test().examples(MathSqrt)
172    }
173}
nu_command/math/sqrt.rs

nu_command/math/
sqrt.rs
