1use nu_engine::command_prelude::*;
2use nu_protocol::{ListStream, shell_error::io::IoError};
3use std::{
4 io::{BufRead, Cursor, ErrorKind},
5 num::NonZeroUsize,
6};
7
8#[derive(Clone)]
9pub struct Chunks;
10
11impl Command for Chunks {
12 fn name(&self) -> &str {
13 "chunks"
14 }
15
16 fn signature(&self) -> Signature {
17 Signature::build("chunks")
18 .input_output_types(vec![
19 (Type::table(), Type::list(Type::table())),
20 (Type::list(Type::Any), Type::list(Type::list(Type::Any))),
21 (Type::Binary, Type::list(Type::Binary)),
22 ])
23 .required(
24 "chunk_size",
25 SyntaxShape::OneOf(vec![SyntaxShape::Int, SyntaxShape::Filesize]),
26 "The size of each chunk.",
27 )
28 .category(Category::Filters)
29 }
30
31 fn description(&self) -> &str {
32 "Divide a list, table or binary input into chunks of `chunk_size`. For binary input, `chunk_size` can also be specified as a filesize."
33 }
34
35 fn extra_description(&self) -> &str {
36 "This command will error if `chunk_size` is negative or zero."
37 }
38
39 fn search_terms(&self) -> Vec<&str> {
40 vec!["batch", "group", "split", "bytes"]
41 }
42
43 fn examples(&self) -> Vec<Example<'_>> {
44 vec![
45 Example {
46 example: "[1 2 3 4] | chunks 2",
47 description: "Chunk a list into pairs",
48 result: Some(Value::test_list(vec![
49 Value::test_list(vec![Value::test_int(1), Value::test_int(2)]),
50 Value::test_list(vec![Value::test_int(3), Value::test_int(4)]),
51 ])),
52 },
53 Example {
54 example: "[[foo bar]; [0 1] [2 3] [4 5] [6 7] [8 9]] | chunks 3",
55 description: "Chunk the rows of a table into triplets",
56 result: Some(Value::test_list(vec![
57 Value::test_list(vec![
58 Value::test_record(record! {
59 "foo" => Value::test_int(0),
60 "bar" => Value::test_int(1),
61 }),
62 Value::test_record(record! {
63 "foo" => Value::test_int(2),
64 "bar" => Value::test_int(3),
65 }),
66 Value::test_record(record! {
67 "foo" => Value::test_int(4),
68 "bar" => Value::test_int(5),
69 }),
70 ]),
71 Value::test_list(vec![
72 Value::test_record(record! {
73 "foo" => Value::test_int(6),
74 "bar" => Value::test_int(7),
75 }),
76 Value::test_record(record! {
77 "foo" => Value::test_int(8),
78 "bar" => Value::test_int(9),
79 }),
80 ]),
81 ])),
82 },
83 Example {
84 example: "0x[11 22 33 44 55 66 77 88] | chunks 3",
85 description: "Chunk the bytes of a binary into triplets",
86 result: Some(Value::test_list(vec![
87 Value::test_binary(vec![0x11, 0x22, 0x33]),
88 Value::test_binary(vec![0x44, 0x55, 0x66]),
89 Value::test_binary(vec![0x77, 0x88]),
90 ])),
91 },
92 Example {
93 example: "open --raw foo.bin | chunks 8kib",
94 description: "Open a binary file and make 8 kibibyte chunks",
95 result: None,
96 },
97 ]
98 }
99
100 fn run(
101 &self,
102 engine_state: &EngineState,
103 stack: &mut Stack,
104 call: &Call,
105 input: PipelineData,
106 ) -> Result<PipelineData, ShellError> {
107 let input = input.into_stream_or_original(engine_state);
108 let head = call.head;
109 let chunk_size: Value = call.req(engine_state, stack, 0)?;
110
111 let size = match chunk_size {
112 Value::Int { val, .. } => {
113 usize::try_from(val).map_err(|_| ShellError::NeedsPositiveValue {
114 span: chunk_size.span(),
115 })
116 }
117 Value::Filesize { val, .. } => {
118 usize::try_from(val).map_err(|_| ShellError::NeedsPositiveValue {
119 span: chunk_size.span(),
120 })
121 }
122 ref val => Err(ShellError::RuntimeTypeMismatch {
123 expected: Type::custom("int or filesize"),
124 actual: val.get_type(),
125 span: val.span(),
126 }),
127 }?;
128
129 let size = NonZeroUsize::try_from(size).map_err(|_| ShellError::IncorrectValue {
130 msg: "`chunk_size` cannot be zero".into(),
131 val_span: chunk_size.span(),
132 call_span: head,
133 })?;
134
135 let is_filesize = matches!(chunk_size, Value::Filesize { .. });
136
137 chunks(engine_state, input, size, head, is_filesize)
138 }
139}
140
141pub fn chunks(
142 engine_state: &EngineState,
143 input: PipelineData,
144 chunk_size: NonZeroUsize,
145 span: Span,
146 is_filesize: bool,
147) -> Result<PipelineData, ShellError> {
148 let from_io_error = IoError::factory(span, None);
149 match input {
150 PipelineData::Value(Value::List { .. }, _) if is_filesize => {
151 Err(ShellError::IncompatibleParametersSingle {
152 msg: "Filesize as chunk size is only supported for binary/byte stream input".into(),
153 span,
154 })
155 }
156 PipelineData::ListStream(_, _) if is_filesize => {
157 Err(ShellError::IncompatibleParametersSingle {
158 msg: "Filesize as chunk size is only supported for binary/byte stream input".into(),
159 span,
160 })
161 }
162 PipelineData::Value(Value::List { vals, .. }, metadata) => {
163 let chunks = ChunksIter::new(vals, chunk_size, span);
164 let stream = ListStream::new(chunks, span, engine_state.signals().clone());
165 Ok(PipelineData::list_stream(stream, metadata))
166 }
167 PipelineData::ListStream(stream, metadata) => {
168 let stream = stream.modify(|iter| ChunksIter::new(iter, chunk_size, span));
169 Ok(PipelineData::list_stream(stream, metadata))
170 }
171 PipelineData::Value(Value::Binary { val, .. }, metadata) => {
172 let chunk_read = ChunkRead {
173 reader: Cursor::new(val),
174 size: chunk_size,
175 };
176 let value_stream = chunk_read.map(move |chunk| match chunk {
177 Ok(chunk) => Value::binary(chunk, span),
178 Err(e) => Value::error(from_io_error(e).into(), span),
179 });
180 let pipeline_data_with_metadata = value_stream.into_pipeline_data_with_metadata(
181 span,
182 engine_state.signals().clone(),
183 metadata,
184 );
185 Ok(pipeline_data_with_metadata)
186 }
187 PipelineData::ByteStream(stream, metadata) => {
188 let pipeline_data = match stream.reader() {
189 None => PipelineData::empty(),
190 Some(reader) => {
191 let chunk_read = ChunkRead {
192 reader,
193 size: chunk_size,
194 };
195 let value_stream = chunk_read.map(move |chunk| match chunk {
196 Ok(chunk) => Value::binary(chunk, span),
197 Err(e) => Value::error(from_io_error(e).into(), span),
198 });
199 value_stream.into_pipeline_data_with_metadata(
200 span,
201 engine_state.signals().clone(),
202 metadata,
203 )
204 }
205 };
206 Ok(pipeline_data)
207 }
208 input => Err(input.unsupported_input_error("list", span)),
209 }
210}
211
212struct ChunksIter<I: Iterator<Item = Value>> {
213 iter: I,
214 size: usize,
215 span: Span,
216}
217
218impl<I: Iterator<Item = Value>> ChunksIter<I> {
219 fn new(iter: impl IntoIterator<IntoIter = I>, size: NonZeroUsize, span: Span) -> Self {
220 Self {
221 iter: iter.into_iter(),
222 size: size.into(),
223 span,
224 }
225 }
226}
227
228impl<I: Iterator<Item = Value>> Iterator for ChunksIter<I> {
229 type Item = Value;
230
231 fn next(&mut self) -> Option<Self::Item> {
232 let first = self.iter.next()?;
233 let mut chunk = Vec::with_capacity(self.size); chunk.push(first);
235 chunk.extend((&mut self.iter).take(self.size - 1));
236 Some(Value::list(chunk, self.span))
237 }
238}
239
240struct ChunkRead<R: BufRead> {
241 reader: R,
242 size: NonZeroUsize,
243}
244
245impl<R: BufRead> Iterator for ChunkRead<R> {
246 type Item = Result<Vec<u8>, std::io::Error>;
247
248 fn next(&mut self) -> Option<Self::Item> {
249 let mut buf = Vec::with_capacity(self.size.get());
250 while buf.len() < self.size.get() {
251 let available = match self.reader.fill_buf() {
252 Ok([]) if buf.is_empty() => return None,
253 Ok([]) => return Some(Ok(buf)),
254 Ok(n) => n,
255 Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
256 Err(e) => return Some(Err(e)),
257 };
258 let needed = self.size.get() - buf.len();
259 let have = available.len().min(needed);
260 buf.extend_from_slice(&available[..have]);
261 self.reader.consume(have);
262 }
263 Some(Ok(buf))
264 }
265}
266
267#[cfg(test)]
268mod test {
269 use std::io::Read;
270
271 use super::*;
272
273 #[test]
274 fn chunk_read() {
275 let s = "hello world";
276 let data = Cursor::new(s);
277 let chunk_read = ChunkRead {
278 reader: data,
279 size: NonZeroUsize::new(4).unwrap(),
280 };
281 let chunks = chunk_read.map(|e| e.unwrap()).collect::<Vec<_>>();
282 assert_eq!(
283 chunks,
284 [&s.as_bytes()[..4], &s.as_bytes()[4..8], &s.as_bytes()[8..]]
285 );
286 }
287
288 #[test]
289 fn chunk_read_stream() {
290 let s = "hello world";
291 let data = Cursor::new(&s[..3])
292 .chain(Cursor::new(&s[3..9]))
293 .chain(Cursor::new(&s[9..]));
294 let chunk_read = ChunkRead {
295 reader: data,
296 size: NonZeroUsize::new(4).unwrap(),
297 };
298 let chunks = chunk_read.map(|e| e.unwrap()).collect::<Vec<_>>();
299 assert_eq!(
300 chunks,
301 [&s.as_bytes()[..4], &s.as_bytes()[4..8], &s.as_bytes()[8..]]
302 );
303 }
304
305 #[test]
306 fn test_examples() -> nu_test_support::Result {
307 nu_test_support::test().examples(Chunks)
308 }
309}