1use std::ops::Bound;
2
3use nu_cmd_base::input_handler::{CmdArgument, operate};
4use nu_engine::command_prelude::*;
5use nu_protocol::{IntRange, Range};
6
7#[derive(Clone)]
8pub struct BytesAt;
9
10struct Arguments {
11 range: IntRange,
12 cell_paths: Option<Vec<CellPath>>,
13}
14
15impl CmdArgument for Arguments {
16 fn take_cell_paths(&mut self) -> Option<Vec<CellPath>> {
17 self.cell_paths.take()
18 }
19}
20
21impl Command for BytesAt {
22 fn name(&self) -> &str {
23 "bytes at"
24 }
25
26 fn signature(&self) -> Signature {
27 Signature::build("bytes at")
28 .input_output_types(vec![
29 (Type::Binary, Type::Binary),
30 (
31 Type::List(Box::new(Type::Binary)),
32 Type::List(Box::new(Type::Binary)),
33 ),
34 (Type::table(), Type::table()),
35 (Type::record(), Type::record()),
36 ])
37 .allow_variants_without_examples(true)
38 .required("range", SyntaxShape::Range, "The range to get bytes.")
39 .rest(
40 "rest",
41 SyntaxShape::CellPath,
42 "For a data structure input, get bytes from data at the given cell paths.",
43 )
44 .category(Category::Bytes)
45 }
46
47 fn description(&self) -> &str {
48 "Get bytes from the input defined by a range."
49 }
50
51 fn search_terms(&self) -> Vec<&str> {
52 vec!["slice"]
53 }
54
55 fn run(
56 &self,
57 engine_state: &EngineState,
58 stack: &mut Stack,
59 call: &Call,
60 input: PipelineData,
61 ) -> Result<PipelineData, ShellError> {
62 let range = match call.req(engine_state, stack, 0)? {
63 Range::IntRange(range) => range,
64 _ => {
65 return Err(ShellError::UnsupportedInput {
66 msg: "Float ranges are not supported for byte streams".into(),
67 input: "value originates from here".into(),
68 msg_span: call.head,
69 input_span: call.head,
70 });
71 }
72 };
73
74 let cell_paths: Vec<CellPath> = call.rest(engine_state, stack, 1)?;
75 let cell_paths = (!cell_paths.is_empty()).then_some(cell_paths);
76
77 if let PipelineData::ByteStream(stream, metadata) = input {
78 let stream = stream.slice(call.head, call.arguments_span(), range)?;
79 let metadata = metadata.map(|m| m.with_content_type(None));
81 Ok(PipelineData::byte_stream(stream, metadata))
82 } else {
83 operate(
84 map_value,
85 Arguments { range, cell_paths },
86 input,
87 call.head,
88 engine_state.signals(),
89 )
90 .map(|mut pipeline| {
91 if let Some(metadata) = pipeline.metadata_mut() {
92 metadata.content_type = None;
94 }
95 pipeline
96 })
97 }
98 }
99
100 fn examples(&self) -> Vec<Example<'_>> {
101 vec![
102 Example {
103 description: "Extract bytes starting from a specific index.",
104 example: "{ data: 0x[33 44 55 10 01 13 10] } | bytes at 3.. data",
105 result: Some(Value::test_record(record! {
106 "data" => Value::test_binary(vec![0x10, 0x01, 0x13, 0x10]),
107 })),
108 },
109 Example {
110 description: "Slice out `0x[10 01 13]` from `0x[33 44 55 10 01 13]`.",
111 example: "0x[33 44 55 10 01 13] | bytes at 3..5",
112 result: Some(Value::test_binary(vec![0x10, 0x01, 0x13])),
113 },
114 Example {
115 description: "Extract bytes from the start up to a specific index.",
116 example: "0x[33 44 55 10 01 13 10] | bytes at ..4",
117 result: Some(Value::test_binary(vec![0x33, 0x44, 0x55, 0x10, 0x01])),
118 },
119 Example {
120 description: "Extract byte `0x[10]` using an exclusive end index.",
121 example: "0x[33 44 55 10 01 13 10] | bytes at 3..<4",
122 result: Some(Value::test_binary(vec![0x10])),
123 },
124 Example {
125 description: "Extract bytes up to a negative index (inclusive).",
126 example: "0x[33 44 55 10 01 13 10] | bytes at ..-4",
127 result: Some(Value::test_binary(vec![0x33, 0x44, 0x55, 0x10])),
128 },
129 Example {
130 description: "Slice bytes across multiple table columns.",
131 example: "[[ColA ColB ColC]; [0x[11 12 13] 0x[14 15 16] 0x[17 18 19]]] | bytes at 1.. ColB ColC",
132 result: Some(Value::test_list(vec![Value::test_record(record! {
133 "ColA" => Value::test_binary(vec![0x11, 0x12, 0x13]),
134 "ColB" => Value::test_binary(vec![0x15, 0x16]),
135 "ColC" => Value::test_binary(vec![0x18, 0x19]),
136 })])),
137 },
138 Example {
139 description: "Extract the last three bytes using a negative start index.",
140 example: "0x[33 44 55 10 01 13 10] | bytes at (-3)..",
141 result: Some(Value::test_binary(vec![0x01, 0x13, 0x10])),
142 },
143 ]
144 }
145}
146
147fn map_value(input: &Value, args: &Arguments, head: Span) -> Value {
148 let range = &args.range;
149 match input {
150 Value::Binary { val, .. } => {
151 let len = val.len() as u64;
152 let start: u64 = range.absolute_start(len);
153 let _start: usize = match start.try_into() {
154 Ok(start) => start,
155 Err(_) => {
156 let span = input.span();
157 return Value::error(
158 ShellError::UnsupportedInput {
159 msg: format!(
160 "Absolute start position {start} was too large for your system arch."
161 ),
162 input: args.range.to_string(),
163 msg_span: span,
164 input_span: span,
165 },
166 head,
167 );
168 }
169 };
170
171 let (start, end) = range.absolute_bounds(val.len());
172 let bytes: Vec<u8> = match end {
173 Bound::Unbounded => val[start..].into(),
174 Bound::Included(end) => val[start..=end].into(),
175 Bound::Excluded(end) => val[start..end].into(),
176 };
177
178 Value::binary(bytes, head)
179 }
180 Value::Error { .. } => input.clone(),
181 other => Value::error(
182 ShellError::UnsupportedInput {
183 msg: "Only binary values are supported".into(),
184 input: format!("input type: {:?}", other.get_type()),
185 msg_span: head,
186 input_span: other.span(),
187 },
188 head,
189 ),
190 }
191}
192
193#[cfg(test)]
194mod test {
195 use super::*;
196
197 #[test]
198 fn test_examples() -> nu_test_support::Result {
199 nu_test_support::test().examples(BytesAt)
200 }
201}