nu_command/math/
sqrt.rs

1use crate::math::utils::ensure_bounded;
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            ])
22            .allow_variants_without_examples(true)
23            .category(Category::Math)
24    }
25
26    fn description(&self) -> &str {
27        "Returns the square root of the input number."
28    }
29
30    fn search_terms(&self) -> Vec<&str> {
31        vec!["square", "root"]
32    }
33
34    fn is_const(&self) -> bool {
35        true
36    }
37
38    fn run(
39        &self,
40        engine_state: &EngineState,
41        _stack: &mut Stack,
42        call: &Call,
43        input: PipelineData,
44    ) -> Result<PipelineData, ShellError> {
45        let head = call.head;
46        // This doesn't match explicit nulls
47        if let PipelineData::Empty = input {
48            return Err(ShellError::PipelineEmpty { dst_span: head });
49        }
50        if let PipelineData::Value(ref v @ Value::Range { ref val, .. }, ..) = input {
51            let span = v.span();
52            ensure_bounded(val, span, head)?;
53        }
54        input.map(move |value| operate(value, head), engine_state.signals())
55    }
56
57    fn run_const(
58        &self,
59        working_set: &StateWorkingSet,
60        call: &Call,
61        input: PipelineData,
62    ) -> Result<PipelineData, ShellError> {
63        let head = call.head;
64        // This doesn't match explicit nulls
65        if let PipelineData::Empty = input {
66            return Err(ShellError::PipelineEmpty { dst_span: head });
67        }
68        if let PipelineData::Value(ref v @ Value::Range { ref val, .. }, ..) = input {
69            let span = v.span();
70            ensure_bounded(val, span, head)?;
71        }
72        input.map(
73            move |value| operate(value, head),
74            working_set.permanent().signals(),
75        )
76    }
77
78    fn examples(&self) -> Vec<Example<'_>> {
79        vec![Example {
80            description: "Compute the square root of each number in a list",
81            example: "[9 16] | math sqrt",
82            result: Some(Value::list(
83                vec![Value::test_float(3.0), Value::test_float(4.0)],
84                Span::test_data(),
85            )),
86        }]
87    }
88}
89
90fn operate(value: Value, head: Span) -> Value {
91    let span = value.span();
92    match value {
93        Value::Int { val, .. } => {
94            let squared = (val as f64).sqrt();
95            if squared.is_nan() {
96                return error_negative_sqrt(head, span);
97            }
98            Value::float(squared, span)
99        }
100        Value::Float { val, .. } => {
101            let squared = val.sqrt();
102            if squared.is_nan() {
103                return error_negative_sqrt(head, span);
104            }
105            Value::float(squared, span)
106        }
107        Value::Error { .. } => value,
108        other => Value::error(
109            ShellError::OnlySupportsThisInputType {
110                exp_input_type: "numeric".into(),
111                wrong_type: other.get_type().to_string(),
112                dst_span: head,
113                src_span: other.span(),
114            },
115            head,
116        ),
117    }
118}
119
120fn error_negative_sqrt(head: Span, span: Span) -> Value {
121    Value::error(
122        ShellError::UnsupportedInput {
123            msg: String::from("Can't square root a negative number"),
124            input: "value originates from here".into(),
125            msg_span: head,
126            input_span: span,
127        },
128        span,
129    )
130}
131
132#[cfg(test)]
133mod test {
134    use super::*;
135
136    #[test]
137    fn test_examples() {
138        use crate::test_examples;
139
140        test_examples(MathSqrt {})
141    }
142}
nu_command/math/sqrt.rs

nu_command/math/
sqrt.rs
