nu_command/bytes/
add.rs

1use nu_cmd_base::input_handler::{operate, CmdArgument};
2use nu_engine::command_prelude::*;
3
4struct Arguments {
5    added_data: Vec<u8>,
6    index: Option<usize>,
7    end: bool,
8    cell_paths: Option<Vec<CellPath>>,
9}
10
11impl CmdArgument for Arguments {
12    fn take_cell_paths(&mut self) -> Option<Vec<CellPath>> {
13        self.cell_paths.take()
14    }
15}
16
17#[derive(Clone)]
18
19pub struct BytesAdd;
20
21impl Command for BytesAdd {
22    fn name(&self) -> &str {
23        "bytes add"
24    }
25
26    fn signature(&self) -> Signature {
27        Signature::build("bytes add")
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("data", SyntaxShape::Binary, "The binary to add.")
39            .named(
40                "index",
41                SyntaxShape::Int,
42                "index to insert binary data",
43                Some('i'),
44            )
45            .switch("end", "add to the end of binary", Some('e'))
46            .rest(
47                "rest",
48                SyntaxShape::CellPath,
49                "For a data structure input, add bytes to the data at the given cell paths.",
50            )
51            .category(Category::Bytes)
52    }
53
54    fn description(&self) -> &str {
55        "Add specified bytes to the input."
56    }
57
58    fn search_terms(&self) -> Vec<&str> {
59        vec!["append", "truncate", "padding"]
60    }
61
62    fn run(
63        &self,
64        engine_state: &EngineState,
65        stack: &mut Stack,
66        call: &Call,
67        input: PipelineData,
68    ) -> Result<PipelineData, ShellError> {
69        let added_data: Vec<u8> = call.req(engine_state, stack, 0)?;
70        let cell_paths: Vec<CellPath> = call.rest(engine_state, stack, 1)?;
71        let cell_paths = (!cell_paths.is_empty()).then_some(cell_paths);
72        let index: Option<usize> = call.get_flag(engine_state, stack, "index")?;
73        let end = call.has_flag(engine_state, stack, "end")?;
74
75        let arg = Arguments {
76            added_data,
77            index,
78            end,
79            cell_paths,
80        };
81        operate(add, arg, input, call.head, engine_state.signals())
82    }
83
84    fn examples(&self) -> Vec<Example> {
85        vec![
86            Example {
87                description: "Add bytes `0x[AA]` to `0x[1F FF AA AA]`",
88                example: "0x[1F FF AA AA] | bytes add 0x[AA]",
89                result: Some(Value::binary(vec![0xAA, 0x1F, 0xFF, 0xAA, 0xAA],
90                    Span::test_data(),
91                )),
92            },
93            Example {
94                description: "Add bytes `0x[AA BB]` to `0x[1F FF AA AA]` at index 1",
95                example: "0x[1F FF AA AA] | bytes add 0x[AA BB] --index 1",
96                result: Some(Value::binary(vec![0x1F, 0xAA, 0xBB, 0xFF, 0xAA, 0xAA],
97                    Span::test_data(),
98                )),
99            },
100            Example {
101                description: "Add bytes `0x[11]` to `0x[FF AA AA]` at the end",
102                example: "0x[FF AA AA] | bytes add 0x[11] --end",
103                result: Some(Value::binary(vec![0xFF, 0xAA, 0xAA, 0x11],
104                    Span::test_data(),
105                )),
106            },
107            Example {
108                description: "Add bytes `0x[11 22 33]` to `0x[FF AA AA]` at the end, at index 1(the index is start from end)",
109                example: "0x[FF AA BB] | bytes add 0x[11 22 33] --end --index 1",
110                result: Some(Value::binary(vec![0xFF, 0xAA, 0x11, 0x22, 0x33, 0xBB],
111                    Span::test_data(),
112                )),
113            },
114        ]
115    }
116}
117
118fn add(val: &Value, args: &Arguments, span: Span) -> Value {
119    let val_span = val.span();
120    match val {
121        Value::Binary { val, .. } => add_impl(val, args, val_span),
122        // Propagate errors by explicitly matching them before the final case.
123        Value::Error { .. } => val.clone(),
124        other => Value::error(
125            ShellError::OnlySupportsThisInputType {
126                exp_input_type: "binary".into(),
127                wrong_type: other.get_type().to_string(),
128                dst_span: span,
129                src_span: other.span(),
130            },
131            span,
132        ),
133    }
134}
135
136fn add_impl(input: &[u8], args: &Arguments, span: Span) -> Value {
137    match args.index {
138        None => {
139            if args.end {
140                let mut added_data = args.added_data.clone();
141                let mut result = input.to_vec();
142                result.append(&mut added_data);
143                Value::binary(result, span)
144            } else {
145                let mut result = args.added_data.clone();
146                let mut input = input.to_vec();
147                result.append(&mut input);
148                Value::binary(result, span)
149            }
150        }
151        Some(mut indx) => {
152            let inserted_index = if args.end {
153                input.len().saturating_sub(indx)
154            } else {
155                if indx > input.len() {
156                    indx = input.len()
157                }
158                indx
159            };
160            let mut result = vec![];
161            let mut prev_data = input[..inserted_index].to_vec();
162            result.append(&mut prev_data);
163            let mut added_data = args.added_data.clone();
164            result.append(&mut added_data);
165            let mut after_data = input[inserted_index..].to_vec();
166            result.append(&mut after_data);
167            Value::binary(result, span)
168        }
169    }
170}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175
176    #[test]
177    fn test_examples() {
178        use crate::test_examples;
179
180        test_examples(BytesAdd {})
181    }
182}
nu_command/bytes/add.rs

nu_command/bytes/
add.rs
