nu_command/filters/
window.rs

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        let window_size: Value = call.req(engine_state, stack, 0)?;
87        let stride: Option<Value> = call.get_flag(engine_state, stack, "stride")?;
88        let remainder = call.has_flag(engine_state, stack, "remainder")?;
89
90        let size =
91            usize::try_from(window_size.as_int()?).map_err(|_| ShellError::NeedsPositiveValue {
92                span: window_size.span(),
93            })?;
94
95        let size = NonZeroUsize::try_from(size).map_err(|_| ShellError::IncorrectValue {
96            msg: "`window_size` cannot be zero".into(),
97            val_span: window_size.span(),
98            call_span: head,
99        })?;
100
101        let stride = if let Some(stride_val) = stride {
102            let stride = usize::try_from(stride_val.as_int()?).map_err(|_| {
103                ShellError::NeedsPositiveValue {
104                    span: stride_val.span(),
105                }
106            })?;
107
108            NonZeroUsize::try_from(stride).map_err(|_| ShellError::IncorrectValue {
109                msg: "`stride` cannot be zero".into(),
110                val_span: stride_val.span(),
111                call_span: head,
112            })?
113        } else {
114            NonZeroUsize::MIN
115        };
116
117        if remainder && size == stride {
118            super::chunks::chunks(engine_state, input, size, head)
119        } else if stride >= size {
120            match input {
121                PipelineData::Value(Value::List { vals, .. }, metadata) => {
122                    let chunks = WindowGapIter::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| WindowGapIter::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        } else {
134            match input {
135                PipelineData::Value(Value::List { vals, .. }, metadata) => {
136                    let chunks = WindowOverlapIter::new(vals, size, stride, remainder, head);
137                    let stream = ListStream::new(chunks, head, engine_state.signals().clone());
138                    Ok(PipelineData::list_stream(stream, metadata))
139                }
140                PipelineData::ListStream(stream, metadata) => {
141                    let stream = stream
142                        .modify(|iter| WindowOverlapIter::new(iter, size, stride, remainder, head));
143                    Ok(PipelineData::list_stream(stream, metadata))
144                }
145                input => Err(input.unsupported_input_error("list", head)),
146            }
147        }
148    }
149}
150
151struct WindowOverlapIter<I: Iterator<Item = Value>> {
152    iter: I,
153    window: Vec<Value>,
154    stride: usize,
155    remainder: bool,
156    span: Span,
157}
158
159impl<I: Iterator<Item = Value>> WindowOverlapIter<I> {
160    fn new(
161        iter: impl IntoIterator<IntoIter = I>,
162        size: NonZeroUsize,
163        stride: NonZeroUsize,
164        remainder: bool,
165        span: Span,
166    ) -> Self {
167        Self {
168            iter: iter.into_iter(),
169            window: Vec::with_capacity(size.into()),
170            stride: stride.into(),
171            remainder,
172            span,
173        }
174    }
175}
176
177impl<I: Iterator<Item = Value>> Iterator for WindowOverlapIter<I> {
178    type Item = Value;
179
180    fn next(&mut self) -> Option<Self::Item> {
181        let len = if self.window.is_empty() {
182            self.window.capacity()
183        } else {
184            self.stride
185        };
186
187        self.window.extend((&mut self.iter).take(len));
188
189        if self.window.len() == self.window.capacity()
190            || (self.remainder && !self.window.is_empty())
191        {
192            let mut next = Vec::with_capacity(self.window.len());
193            next.extend(self.window.iter().skip(self.stride).cloned());
194            let window = std::mem::replace(&mut self.window, next);
195            Some(Value::list(window, self.span))
196        } else {
197            None
198        }
199    }
200}
201
202struct WindowGapIter<I: Iterator<Item = Value>> {
203    iter: I,
204    size: usize,
205    skip: usize,
206    first: bool,
207    remainder: bool,
208    span: Span,
209}
210
211impl<I: Iterator<Item = Value>> WindowGapIter<I> {
212    fn new(
213        iter: impl IntoIterator<IntoIter = I>,
214        size: NonZeroUsize,
215        stride: NonZeroUsize,
216        remainder: bool,
217        span: Span,
218    ) -> Self {
219        let size = size.into();
220        Self {
221            iter: iter.into_iter(),
222            size,
223            skip: stride.get() - size,
224            first: true,
225            remainder,
226            span,
227        }
228    }
229}
230
231impl<I: Iterator<Item = Value>> Iterator for WindowGapIter<I> {
232    type Item = Value;
233
234    fn next(&mut self) -> Option<Self::Item> {
235        let mut window = Vec::with_capacity(self.size);
236        window.extend(
237            (&mut self.iter)
238                .skip(if self.first { 0 } else { self.skip })
239                .take(self.size),
240        );
241
242        self.first = false;
243
244        if window.len() == self.size || (self.remainder && !window.is_empty()) {
245            Some(Value::list(window, self.span))
246        } else {
247            None
248        }
249    }
250}
251
252#[cfg(test)]
253mod test {
254    use super::*;
255
256    #[test]
257    fn test_examples() -> nu_test_support::Result {
258        nu_test_support::test().examples(Window)
259    }
260}
nu_command/filters/window.rs

nu_command/filters/
window.rs
