1use nu_engine::command_prelude::*;
2use nu_protocol::{PipelineIterator, Range};
3use std::collections::VecDeque;
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 .rest(
22 "rest",
23 SyntaxShape::Any,
24 "The row numbers or ranges to drop.",
25 )
26 .category(Category::Filters)
27 }
28
29 fn description(&self) -> &str {
30 "Drop the selected rows."
31 }
32
33 fn search_terms(&self) -> Vec<&str> {
34 vec!["delete", "remove", "index"]
35 }
36
37 fn examples(&self) -> Vec<Example<'_>> {
38 vec![
39 Example {
40 example: "[sam,sarah,2,3,4,5] | drop nth 0 1 2",
41 description: "Drop the first, second, and third row",
42 result: Some(Value::list(
43 vec![Value::test_int(3), Value::test_int(4), Value::test_int(5)],
44 Span::test_data(),
45 )),
46 },
47 Example {
48 example: "[0,1,2,3,4,5] | drop nth 0 1 2",
49 description: "Drop the first, second, and third row",
50 result: Some(Value::list(
51 vec![Value::test_int(3), Value::test_int(4), Value::test_int(5)],
52 Span::test_data(),
53 )),
54 },
55 Example {
56 example: "[0,1,2,3,4,5] | drop nth 0 2 4",
57 description: "Drop rows 0 2 4",
58 result: Some(Value::list(
59 vec![Value::test_int(1), Value::test_int(3), Value::test_int(5)],
60 Span::test_data(),
61 )),
62 },
63 Example {
64 example: "[0,1,2,3,4,5] | drop nth 2 0 4",
65 description: "Drop rows 2 0 4",
66 result: Some(Value::list(
67 vec![Value::test_int(1), Value::test_int(3), Value::test_int(5)],
68 Span::test_data(),
69 )),
70 },
71 Example {
72 description: "Drop range rows from second to fourth",
73 example: "[first second third fourth fifth] | drop nth (1..3)",
74 result: Some(Value::list(
75 vec![Value::test_string("first"), Value::test_string("fifth")],
76 Span::test_data(),
77 )),
78 },
79 Example {
80 example: "[0,1,2,3,4,5] | drop nth 1..",
81 description: "Drop all rows except first row",
82 result: Some(Value::list(vec![Value::test_int(0)], Span::test_data())),
83 },
84 Example {
85 example: "[0,1,2,3,4,5] | drop nth 3..",
86 description: "Drop rows 3,4,5",
87 result: Some(Value::list(
88 vec![Value::test_int(0), Value::test_int(1), Value::test_int(2)],
89 Span::test_data(),
90 )),
91 },
92 ]
93 }
94
95 fn run(
96 &self,
97 engine_state: &EngineState,
98 stack: &mut Stack,
99 call: &Call,
100 input: PipelineData,
101 ) -> Result<PipelineData, ShellError> {
102 let head = call.head;
103 let metadata = input.metadata();
104
105 let args: Vec<Value> = call.rest(engine_state, stack, 0)?;
106 if args.is_empty() {
107 return Ok(input);
108 }
109
110 let (rows_to_drop, min_unbounded_start) = get_rows_to_drop(&args, head)?;
111
112 let input = if let Some(cutoff) = min_unbounded_start {
113 input
114 .into_iter()
115 .take(cutoff)
116 .into_pipeline_data_with_metadata(
117 head,
118 engine_state.signals().clone(),
119 metadata.clone(),
120 )
121 } else {
122 input
123 };
124
125 Ok(DropNthIterator {
126 input: input.into_iter(),
127 rows: rows_to_drop,
128 current: 0,
129 }
130 .into_pipeline_data_with_metadata(head, engine_state.signals().clone(), metadata))
131 }
132}
133
134fn get_rows_to_drop(
135 args: &[Value],
136 head: Span,
137) -> Result<(VecDeque<usize>, Option<usize>), ShellError> {
138 let mut rows_to_drop = Vec::new();
139 let mut min_unbounded_start: Option<usize> = None;
140
141 for value in args {
142 if let Ok(i) = value.as_int() {
143 if i < 0 {
144 return Err(ShellError::UnsupportedInput {
145 msg: "drop nth accepts only positive ints".into(),
146 input: "value originates from here".into(),
147 msg_span: head,
148 input_span: value.span(),
149 });
150 }
151 rows_to_drop.push(i as usize);
152 } else if let Ok(range) = value.as_range() {
153 match range {
154 Range::IntRange(range) => {
155 let start = range.start();
156 if start < 0 {
157 return Err(ShellError::UnsupportedInput {
158 msg: "drop nth accepts only positive ints".into(),
159 input: "value originates from here".into(),
160 msg_span: head,
161 input_span: value.span(),
162 });
163 }
164
165 match range.end() {
166 Bound::Included(end) => {
167 if end < start {
168 return Err(ShellError::UnsupportedInput {
169 msg: "The upper bound must be greater than or equal to the lower bound".into(),
170 input: "value originates from here".into(),
171 msg_span: head,
172 input_span: value.span(),
173 });
174 }
175 rows_to_drop.extend((start as usize)..=(end as usize));
176 }
177 Bound::Excluded(end) => {
178 if end <= start {
179 return Err(ShellError::UnsupportedInput {
180 msg: "The upper bound must be greater than the lower bound"
181 .into(),
182 input: "value originates from here".into(),
183 msg_span: head,
184 input_span: value.span(),
185 });
186 }
187 rows_to_drop.extend((start as usize)..(end as usize));
188 }
189 Bound::Unbounded => {
190 let start_usize = start as usize;
191 min_unbounded_start = Some(
192 min_unbounded_start.map_or(start_usize, |s| s.min(start_usize)),
193 );
194 }
195 }
196 }
197 Range::FloatRange(_) => {
198 return Err(ShellError::UnsupportedInput {
199 msg: "float range not supported".into(),
200 input: "value originates from here".into(),
201 msg_span: head,
202 input_span: value.span(),
203 });
204 }
205 }
206 } else {
207 return Err(ShellError::TypeMismatch {
208 err_message: "Expected int or range".into(),
209 span: value.span(),
210 });
211 }
212 }
213
214 rows_to_drop.sort_unstable();
215 rows_to_drop.dedup();
216
217 Ok((VecDeque::from(rows_to_drop), min_unbounded_start))
218}
219
220struct DropNthIterator {
221 input: PipelineIterator,
222 rows: VecDeque<usize>,
223 current: usize,
224}
225
226impl Iterator for DropNthIterator {
227 type Item = Value;
228
229 fn next(&mut self) -> Option<Self::Item> {
230 loop {
231 if let Some(row) = self.rows.front() {
232 if self.current == *row {
233 self.rows.pop_front();
234 self.current += 1;
235 let _ = self.input.next();
236 continue;
237 } else {
238 self.current += 1;
239 return self.input.next();
240 }
241 } else {
242 return self.input.next();
243 }
244 }
245 }
246}
247
248#[cfg(test)]
249mod test {
250 use super::*;
251
252 #[test]
253 fn test_examples() {
254 use crate::test_examples;
255
256 test_examples(DropNth {})
257 }
258}