nu_command/math/
utils.rs

1use core::slice;
2use indexmap::IndexMap;
3use nu_protocol::{
4    IntoPipelineData, PipelineData, Range, ShellError, Signals, Span, Value, engine::Call,
5};
6
7pub fn run_with_function(
8    call: &Call,
9    input: PipelineData,
10    mf: impl Fn(&[Value], Span, Span) -> Result<Value, ShellError>,
11) -> Result<PipelineData, ShellError> {
12    let name = call.head;
13    let res = calculate(input, name, mf);
14    match res {
15        Ok(v) => Ok(v.into_pipeline_data()),
16        Err(e) => Err(e),
17    }
18}
19
20fn helper_for_tables(
21    values: &[Value],
22    val_span: Span,
23    name: Span,
24    mf: impl Fn(&[Value], Span, Span) -> Result<Value, ShellError>,
25) -> Result<Value, ShellError> {
26    // If we are not dealing with Primitives, then perhaps we are dealing with a table
27    // Create a key for each column name
28    let mut column_values = IndexMap::new();
29    for val in values {
30        match val {
31            Value::Record { val, .. } => {
32                for (key, value) in &**val {
33                    column_values
34                        .entry(key.clone())
35                        .and_modify(|v: &mut Vec<Value>| v.push(value.clone()))
36                        .or_insert_with(|| vec![value.clone()]);
37                }
38            }
39            Value::Error { error, .. } => return Err(*error.clone()),
40            _ => {
41                //Turns out we are not dealing with a table
42                return mf(values, val.span(), name);
43            }
44        }
45    }
46    // The mathematical function operates over the columns of the table
47    let mut column_totals = IndexMap::new();
48    for (col_name, col_vals) in column_values {
49        if let Ok(out) = mf(&col_vals, val_span, name) {
50            column_totals.insert(col_name, out);
51        }
52    }
53    if column_totals.keys().len() == 0 {
54        return Err(ShellError::UnsupportedInput {
55            msg: "Unable to give a result with this input".to_string(),
56            input: "value originates from here".into(),
57            msg_span: name,
58            input_span: val_span,
59        });
60    }
61
62    Ok(Value::record(column_totals.into_iter().collect(), name))
63}
64
65pub fn calculate(
66    values: PipelineData,
67    name: Span,
68    mf: impl Fn(&[Value], Span, Span) -> Result<Value, ShellError>,
69) -> Result<Value, ShellError> {
70    // TODO implement spans for ListStream, thus negating the need for unwrap_or().
71    let span = values.span().unwrap_or(name);
72    match values {
73        PipelineData::ListStream(s, ..) => {
74            helper_for_tables(&s.into_iter().collect::<Vec<Value>>(), span, name, mf)
75        }
76        PipelineData::Value(Value::List { ref vals, .. }, ..) => match &vals[..] {
77            [Value::Record { .. }, _end @ ..] => helper_for_tables(
78                vals,
79                values.span().expect("PipelineData::Value had no span"),
80                name,
81                mf,
82            ),
83            _ => mf(vals, span, name),
84        },
85        PipelineData::Value(Value::Record { val, .. }, ..) => {
86            let mut record = val.into_owned();
87            record
88                .iter_mut()
89                .try_for_each(|(_, val)| -> Result<(), ShellError> {
90                    *val = mf(slice::from_ref(val), span, name)?;
91                    Ok(())
92                })?;
93            Ok(Value::record(record, span))
94        }
95        PipelineData::Value(Value::Range { val, .. }, ..) => {
96            ensure_bounded(val.as_ref(), span, name)?;
97            let new_vals: Result<Vec<Value>, ShellError> = val
98                .into_range_iter(span, Signals::empty())
99                .map(|val| mf(&[val], span, name))
100                .collect();
101
102            mf(&new_vals?, span, name)
103        }
104        PipelineData::Value(val, ..) => mf(&[val], span, name),
105        PipelineData::Empty => Err(ShellError::PipelineEmpty { dst_span: name }),
106        val => Err(ShellError::UnsupportedInput {
107            msg: "Only ints, floats, lists, records, or ranges are supported".into(),
108            input: "value originates from here".into(),
109            msg_span: name,
110            input_span: val
111                .span()
112                .expect("non-Empty non-ListStream PipelineData had no span"),
113        }),
114    }
115}
116
117pub fn ensure_bounded(range: &Range, val_span: Span, call_span: Span) -> Result<(), ShellError> {
118    if range.is_bounded() {
119        return Ok(());
120    }
121    Err(ShellError::IncorrectValue {
122        msg: "Range must be bounded".to_string(),
123        val_span,
124        call_span,
125    })
126}
nu_command/math/utils.rs

nu_command/math/
utils.rs
