1use nu_engine::command_prelude::*;
2use nu_protocol::ListStream;
3use std::num::NonZeroUsize;
4
5#[derive(Clone)]
6pub struct Window;
7
8impl Command for Window {
9 fn name(&self) -> &str {
10 "window"
11 }
12
13 fn signature(&self) -> Signature {
14 Signature::build("window")
15 .input_output_types(vec![(
16 Type::list(Type::Any),
17 Type::list(Type::list(Type::Any)),
18 )])
19 .required("window_size", SyntaxShape::Int, "The size of each window.")
20 .named(
21 "stride",
22 SyntaxShape::Int,
23 "The number of rows to slide over between windows.",
24 Some('s'),
25 )
26 .switch(
27 "remainder",
28 "Yield last chunks even if they have fewer elements than size.",
29 Some('r'),
30 )
31 .category(Category::Filters)
32 }
33
34 fn description(&self) -> &str {
35 "Creates a sliding window of `window_size` that slide by n rows/elements across input."
36 }
37
38 fn extra_description(&self) -> &str {
39 "This command will error if `window_size` or `stride` are negative or zero."
40 }
41
42 fn examples(&self) -> Vec<Example<'_>> {
43 vec![
44 Example {
45 example: "[1 2 3 4] | window 2",
46 description: "A sliding window of two elements.",
47 result: Some(Value::test_list(vec![
48 Value::test_list(vec![Value::test_int(1), Value::test_int(2)]),
49 Value::test_list(vec![Value::test_int(2), Value::test_int(3)]),
50 Value::test_list(vec![Value::test_int(3), Value::test_int(4)]),
51 ])),
52 },
53 Example {
54 example: "[1, 2, 3, 4, 5, 6, 7, 8] | window 2 --stride 3",
55 description: "A sliding window of two elements, with a stride of 3.",
56 result: Some(Value::test_list(vec![
57 Value::test_list(vec![Value::test_int(1), Value::test_int(2)]),
58 Value::test_list(vec![Value::test_int(4), Value::test_int(5)]),
59 Value::test_list(vec![Value::test_int(7), Value::test_int(8)]),
60 ])),
61 },
62 Example {
63 example: "[1, 2, 3, 4, 5] | window 3 --stride 3 --remainder",
64 description: "A sliding window of equal stride that includes remainder. Equivalent to chunking.",
65 result: Some(Value::test_list(vec![
66 Value::test_list(vec![
67 Value::test_int(1),
68 Value::test_int(2),
69 Value::test_int(3),
70 ]),
71 Value::test_list(vec![Value::test_int(4), Value::test_int(5)]),
72 ])),
73 },
74 ]
75 }
76
77 fn run(
78 &self,
79 engine_state: &EngineState,
80 stack: &mut Stack,
81 call: &Call,
82 input: PipelineData,
83 ) -> Result<PipelineData, ShellError> {
84 let input = input.into_stream_or_original(engine_state);
85 let head = call.head;
86
87 let fix_call_span = |err: ShellError| match err {
88 ShellError::IncorrectValue { msg, val_span, .. } => ShellError::IncorrectValue {
89 msg,
90 val_span,
91 call_span: call.head,
92 },
93 _ => err,
94 };
95
96 let size: NonZeroUsize = call.req(engine_state, stack, 0).map_err(fix_call_span)?;
97 let stride: NonZeroUsize = call
98 .get_flag(engine_state, stack, "stride")
99 .map_err(fix_call_span)?
100 .unwrap_or(NonZeroUsize::MIN);
101 let remainder = call.has_flag(engine_state, stack, "remainder")?;
102
103 if remainder && size == stride {
104 super::chunks::chunks(engine_state, input, size, head)
105 } else if stride >= size {
106 match input {
107 PipelineData::Value(Value::List { vals, .. }, metadata) => {
108 let chunks = WindowGapIter::new(vals, size, stride, remainder, head);
109 let stream = ListStream::new(chunks, head, engine_state.signals().clone());
110 Ok(PipelineData::list_stream(stream, metadata))
111 }
112 PipelineData::ListStream(stream, metadata) => {
113 let stream = stream
114 .modify(|iter| WindowGapIter::new(iter, size, stride, remainder, head));
115 Ok(PipelineData::list_stream(stream, metadata))
116 }
117 input => Err(input.unsupported_input_error("list", head)),
118 }
119 } else {
120 match input {
121 PipelineData::Value(Value::List { vals, .. }, metadata) => {
122 let chunks = WindowOverlapIter::new(vals, size, stride, remainder, head);
123 let stream = ListStream::new(chunks, head, engine_state.signals().clone());
124 Ok(PipelineData::list_stream(stream, metadata))
125 }
126 PipelineData::ListStream(stream, metadata) => {
127 let stream = stream
128 .modify(|iter| WindowOverlapIter::new(iter, size, stride, remainder, head));
129 Ok(PipelineData::list_stream(stream, metadata))
130 }
131 input => Err(input.unsupported_input_error("list", head)),
132 }
133 }
134 }
135}
136
137struct WindowOverlapIter<I: Iterator<Item = Value>> {
138 iter: I,
139 window: Vec<Value>,
140 stride: usize,
141 remainder: bool,
142 span: Span,
143}
144
145impl<I: Iterator<Item = Value>> WindowOverlapIter<I> {
146 fn new(
147 iter: impl IntoIterator<IntoIter = I>,
148 size: NonZeroUsize,
149 stride: NonZeroUsize,
150 remainder: bool,
151 span: Span,
152 ) -> Self {
153 Self {
154 iter: iter.into_iter(),
155 window: Vec::with_capacity(size.into()),
156 stride: stride.into(),
157 remainder,
158 span,
159 }
160 }
161}
162
163impl<I: Iterator<Item = Value>> Iterator for WindowOverlapIter<I> {
164 type Item = Value;
165
166 fn next(&mut self) -> Option<Self::Item> {
167 let len = if self.window.is_empty() {
168 self.window.capacity()
169 } else {
170 self.stride
171 };
172
173 self.window.extend((&mut self.iter).take(len));
174
175 if self.window.len() == self.window.capacity()
176 || (self.remainder && !self.window.is_empty())
177 {
178 let mut next = Vec::with_capacity(self.window.len());
179 next.extend(self.window.iter().skip(self.stride).cloned());
180 let window = std::mem::replace(&mut self.window, next);
181 Some(Value::list(window, self.span))
182 } else {
183 None
184 }
185 }
186}
187
188struct WindowGapIter<I: Iterator<Item = Value>> {
189 iter: I,
190 size: usize,
191 skip: usize,
192 first: bool,
193 remainder: bool,
194 span: Span,
195}
196
197impl<I: Iterator<Item = Value>> WindowGapIter<I> {
198 fn new(
199 iter: impl IntoIterator<IntoIter = I>,
200 size: NonZeroUsize,
201 stride: NonZeroUsize,
202 remainder: bool,
203 span: Span,
204 ) -> Self {
205 let size = size.into();
206 Self {
207 iter: iter.into_iter(),
208 size,
209 skip: stride.get() - size,
210 first: true,
211 remainder,
212 span,
213 }
214 }
215}
216
217impl<I: Iterator<Item = Value>> Iterator for WindowGapIter<I> {
218 type Item = Value;
219
220 fn next(&mut self) -> Option<Self::Item> {
221 let mut window = Vec::with_capacity(self.size);
222 window.extend(
223 (&mut self.iter)
224 .skip(if self.first { 0 } else { self.skip })
225 .take(self.size),
226 );
227
228 self.first = false;
229
230 if window.len() == self.size || (self.remainder && !window.is_empty()) {
231 Some(Value::list(window, self.span))
232 } else {
233 None
234 }
235 }
236}
237
238#[cfg(test)]
239mod test {
240 use super::*;
241
242 #[test]
243 fn test_examples() -> nu_test_support::Result {
244 nu_test_support::test().examples(Window)
245 }
246}