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

nu_command/filters/
window.rs
