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 mut input: PipelineData,
101 ) -> Result<PipelineData, ShellError> {
102 let head = call.head;
103 let metadata = input.take_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(head, engine_state.signals().clone())
117 } else {
118 input
119 };
120
121 Ok(DropNthIterator {
122 input: input.into_iter(),
123 rows: rows_to_drop,
124 current: 0,
125 }
126 .into_pipeline_data_with_metadata(head, engine_state.signals().clone(), metadata))
127 }
128}
129
130fn get_rows_to_drop(
131 args: &[Value],
132 head: Span,
133) -> Result<(VecDeque<usize>, Option<usize>), ShellError> {
134 let mut rows_to_drop = Vec::new();
135 let mut min_unbounded_start: Option<usize> = None;
136
137 for value in args {
138 if let Ok(i) = value.as_int() {
139 if i < 0 {
140 return Err(ShellError::UnsupportedInput {
141 msg: "drop nth accepts only positive ints".into(),
142 input: "value originates from here".into(),
143 msg_span: head,
144 input_span: value.span(),
145 });
146 }
147 rows_to_drop.push(i as usize);
148 } else if let Ok(range) = value.as_range() {
149 match range {
150 Range::IntRange(range) => {
151 let start = range.start();
152 if start < 0 {
153 return Err(ShellError::UnsupportedInput {
154 msg: "drop nth accepts only positive ints".into(),
155 input: "value originates from here".into(),
156 msg_span: head,
157 input_span: value.span(),
158 });
159 }
160
161 match range.end() {
162 Bound::Included(end) => {
163 if end < start {
164 return Err(ShellError::UnsupportedInput {
165 msg: "The upper bound must be greater than or equal to the lower bound".into(),
166 input: "value originates from here".into(),
167 msg_span: head,
168 input_span: value.span(),
169 });
170 }
171 rows_to_drop.extend((start as usize)..=(end as usize));
172 }
173 Bound::Excluded(end) => {
174 if end <= start {
175 return Err(ShellError::UnsupportedInput {
176 msg: "The upper bound must be greater than the lower bound"
177 .into(),
178 input: "value originates from here".into(),
179 msg_span: head,
180 input_span: value.span(),
181 });
182 }
183 rows_to_drop.extend((start as usize)..(end as usize));
184 }
185 Bound::Unbounded => {
186 let start_usize = start as usize;
187 min_unbounded_start = Some(
188 min_unbounded_start.map_or(start_usize, |s| s.min(start_usize)),
189 );
190 }
191 }
192 }
193 Range::FloatRange(_) => {
194 return Err(ShellError::UnsupportedInput {
195 msg: "float range not supported".into(),
196 input: "value originates from here".into(),
197 msg_span: head,
198 input_span: value.span(),
199 });
200 }
201 }
202 } else {
203 return Err(ShellError::TypeMismatch {
204 err_message: "Expected int or range".into(),
205 span: value.span(),
206 });
207 }
208 }
209
210 rows_to_drop.sort_unstable();
211 rows_to_drop.dedup();
212
213 Ok((VecDeque::from(rows_to_drop), min_unbounded_start))
214}
215
216struct DropNthIterator {
217 input: PipelineIterator,
218 rows: VecDeque<usize>,
219 current: usize,
220}
221
222impl Iterator for DropNthIterator {
223 type Item = Value;
224
225 fn next(&mut self) -> Option<Self::Item> {
226 loop {
227 if let Some(row) = self.rows.front() {
228 if self.current == *row {
229 self.rows.pop_front();
230 self.current += 1;
231 let _ = self.input.next();
232 continue;
233 } else {
234 self.current += 1;
235 return self.input.next();
236 }
237 } else {
238 return self.input.next();
239 }
240 }
241 }
242}
243
244#[cfg(test)]
245mod test {
246 use super::*;
247
248 #[test]
249 fn test_examples() -> nu_test_support::Result {
250 nu_test_support::test().examples(DropNth)
251 }
252}