nu_command/filters/drop/
nth.rs

1use itertools::Either;
2use nu_engine::command_prelude::*;
3use nu_protocol::{PipelineIterator, Range};
4use std::ops::Bound;
5
6#[derive(Clone)]
7pub struct DropNth;
8
9impl Command for DropNth {
10    fn name(&self) -> &str {
11        "drop nth"
12    }
13
14    fn signature(&self) -> Signature {
15        Signature::build("drop nth")
16            .input_output_types(vec![
17                (Type::Range, Type::list(Type::Number)),
18                (Type::list(Type::Any), Type::list(Type::Any)),
19            ])
20            .allow_variants_without_examples(true)
21            .required(
22                "row number or row range",
23                // FIXME: we can make this accept either Int or Range when we can compose SyntaxShapes
24                SyntaxShape::Any,
25                "The number of the row to drop or a range to drop consecutive rows.",
26            )
27            .rest("rest", SyntaxShape::Any, "The number of the row to drop.")
28            .category(Category::Filters)
29    }
30
31    fn description(&self) -> &str {
32        "Drop the selected rows."
33    }
34
35    fn search_terms(&self) -> Vec<&str> {
36        vec!["delete", "remove", "index"]
37    }
38
39    fn examples(&self) -> Vec<Example> {
40        vec![
41            Example {
42                example: "[sam,sarah,2,3,4,5] | drop nth 0 1 2",
43                description: "Drop the first, second, and third row",
44                result: Some(Value::list(
45                    vec![Value::test_int(3), Value::test_int(4), Value::test_int(5)],
46                    Span::test_data(),
47                )),
48            },
49            Example {
50                example: "[0,1,2,3,4,5] | drop nth 0 1 2",
51                description: "Drop the first, second, and third row",
52                result: Some(Value::list(
53                    vec![Value::test_int(3), Value::test_int(4), Value::test_int(5)],
54                    Span::test_data(),
55                )),
56            },
57            Example {
58                example: "[0,1,2,3,4,5] | drop nth 0 2 4",
59                description: "Drop rows 0 2 4",
60                result: Some(Value::list(
61                    vec![Value::test_int(1), Value::test_int(3), Value::test_int(5)],
62                    Span::test_data(),
63                )),
64            },
65            Example {
66                example: "[0,1,2,3,4,5] | drop nth 2 0 4",
67                description: "Drop rows 2 0 4",
68                result: Some(Value::list(
69                    vec![Value::test_int(1), Value::test_int(3), Value::test_int(5)],
70                    Span::test_data(),
71                )),
72            },
73            Example {
74                description: "Drop range rows from second to fourth",
75                example: "[first second third fourth fifth] | drop nth (1..3)",
76                result: Some(Value::list(
77                    vec![Value::test_string("first"), Value::test_string("fifth")],
78                    Span::test_data(),
79                )),
80            },
81            Example {
82                example: "[0,1,2,3,4,5] | drop nth 1..",
83                description: "Drop all rows except first row",
84                result: Some(Value::list(vec![Value::test_int(0)], Span::test_data())),
85            },
86            Example {
87                example: "[0,1,2,3,4,5] | drop nth 3..",
88                description: "Drop rows 3,4,5",
89                result: Some(Value::list(
90                    vec![Value::test_int(0), Value::test_int(1), Value::test_int(2)],
91                    Span::test_data(),
92                )),
93            },
94        ]
95    }
96
97    fn run(
98        &self,
99        engine_state: &EngineState,
100        stack: &mut Stack,
101        call: &Call,
102        input: PipelineData,
103    ) -> Result<PipelineData, ShellError> {
104        let head = call.head;
105        let metadata = input.metadata();
106        let number_or_range = extract_int_or_range(engine_state, stack, call)?;
107
108        let rows = match number_or_range.item {
109            Either::Left(row_number) => {
110                let and_rows: Vec<Spanned<i64>> = call.rest(engine_state, stack, 1)?;
111                let mut rows: Vec<_> = and_rows.into_iter().map(|x| x.item as usize).collect();
112                rows.push(row_number as usize);
113                rows.sort_unstable();
114                rows
115            }
116            Either::Right(Range::FloatRange(_)) => {
117                return Err(ShellError::UnsupportedInput {
118                    msg: "float range".into(),
119                    input: "value originates from here".into(),
120                    msg_span: head,
121                    input_span: number_or_range.span,
122                });
123            }
124            Either::Right(Range::IntRange(range)) => {
125                // check for negative range inputs, e.g., (2..-5)
126                let end_negative = match range.end() {
127                    Bound::Included(end) | Bound::Excluded(end) => end < 0,
128                    Bound::Unbounded => false,
129                };
130                if range.start().is_negative() || end_negative {
131                    return Err(ShellError::UnsupportedInput {
132                        msg: "drop nth accepts only positive ints".into(),
133                        input: "value originates from here".into(),
134                        msg_span: head,
135                        input_span: number_or_range.span,
136                    });
137                }
138                // check if the upper bound is smaller than the lower bound, e.g., do not accept 4..2
139                if range.step() < 0 {
140                    return Err(ShellError::UnsupportedInput {
141                        msg: "The upper bound needs to be equal or larger to the lower bound"
142                            .into(),
143                        input: "value originates from here".into(),
144                        msg_span: head,
145                        input_span: number_or_range.span,
146                    });
147                }
148
149                let start = range.start() as usize;
150
151                let end = match range.end() {
152                    Bound::Included(end) => end as usize,
153                    Bound::Excluded(end) => (end - 1) as usize,
154                    Bound::Unbounded => {
155                        return Ok(input
156                            .into_iter()
157                            .take(start)
158                            .into_pipeline_data_with_metadata(
159                                head,
160                                engine_state.signals().clone(),
161                                metadata,
162                            ));
163                    }
164                };
165
166                let end = if let PipelineData::Value(Value::List { vals, .. }, _) = &input {
167                    end.min(vals.len() - 1)
168                } else {
169                    end
170                };
171
172                (start..=end).collect()
173            }
174        };
175
176        Ok(DropNthIterator {
177            input: input.into_iter(),
178            rows,
179            current: 0,
180        }
181        .into_pipeline_data_with_metadata(head, engine_state.signals().clone(), metadata))
182    }
183}
184
185fn extract_int_or_range(
186    engine_state: &EngineState,
187    stack: &mut Stack,
188    call: &Call,
189) -> Result<Spanned<Either<i64, Range>>, ShellError> {
190    let value: Value = call.req(engine_state, stack, 0)?;
191
192    let int_opt = value.as_int().map(Either::Left).ok();
193    let range_opt = value.as_range().map(Either::Right).ok();
194
195    int_opt
196        .or(range_opt)
197        .ok_or_else(|| ShellError::TypeMismatch {
198            err_message: "int or range".into(),
199            span: value.span(),
200        })
201        .map(|either| Spanned {
202            item: either,
203            span: value.span(),
204        })
205}
206
207struct DropNthIterator {
208    input: PipelineIterator,
209    rows: Vec<usize>,
210    current: usize,
211}
212
213impl Iterator for DropNthIterator {
214    type Item = Value;
215
216    fn next(&mut self) -> Option<Self::Item> {
217        loop {
218            if let Some(row) = self.rows.first() {
219                if self.current == *row {
220                    self.rows.remove(0);
221                    self.current += 1;
222                    let _ = self.input.next();
223                    continue;
224                } else {
225                    self.current += 1;
226                    return self.input.next();
227                }
228            } else {
229                return self.input.next();
230            }
231        }
232    }
233}
234
235#[cfg(test)]
236mod test {
237    use super::*;
238
239    #[test]
240    fn test_examples() {
241        use crate::test_examples;
242
243        test_examples(DropNth {})
244    }
245}
nu_command/filters/drop/nth.rs

nu_command/filters/drop/
nth.rs
