nushell/crates/nu-command/src/bytes/at.rs

use std::ops::Bound;

use nu_cmd_base::input_handler::{operate, CmdArgument};
use nu_engine::command_prelude::*;
use nu_protocol::{IntRange, Range};

#[derive(Clone)]
pub struct BytesAt;

struct Arguments {
    range: IntRange,
    cell_paths: Option<Vec<CellPath>>,
}

impl CmdArgument for Arguments {
    fn take_cell_paths(&mut self) -> Option<Vec<CellPath>> {
        self.cell_paths.take()
    }
}

impl Command for BytesAt {
    fn name(&self) -> &str {
        "bytes at"
    }

    fn signature(&self) -> Signature {
        Signature::build("bytes at")
            .input_output_types(vec![
                (Type::Binary, Type::Binary),
                (
                    Type::List(Box::new(Type::Binary)),
                    Type::List(Box::new(Type::Binary)),
                ),
                (Type::table(), Type::table()),
                (Type::record(), Type::record()),
            ])
            .allow_variants_without_examples(true)
            .required("range", SyntaxShape::Range, "The range to get bytes.")
            .rest(
                "rest",
                SyntaxShape::CellPath,
                "For a data structure input, get bytes from data at the given cell paths.",
            )
            .category(Category::Bytes)
    }

    fn description(&self) -> &str {
        "Get bytes defined by a range."
    }

    fn search_terms(&self) -> Vec<&str> {
        vec!["slice"]
    }

    fn run(
        &self,
        engine_state: &EngineState,
        stack: &mut Stack,
        call: &Call,
        input: PipelineData,
    ) -> Result<PipelineData, ShellError> {
        let range = match call.req(engine_state, stack, 0)? {
            Range::IntRange(range) => range,
            _ => {
                return Err(ShellError::UnsupportedInput {
                    msg: "Float ranges are not supported for byte streams".into(),
                    input: "value originates from here".into(),
                    msg_span: call.head,
                    input_span: call.head,
                })
            }
        };

        let cell_paths: Vec<CellPath> = call.rest(engine_state, stack, 1)?;
        let cell_paths = (!cell_paths.is_empty()).then_some(cell_paths);

        if let PipelineData::ByteStream(stream, metadata) = input {
            let stream = stream.slice(call.head, call.arguments_span(), range)?;
            Ok(PipelineData::ByteStream(stream, metadata))
        } else {
            operate(
                map_value,
                Arguments { range, cell_paths },
                input,
                call.head,
                engine_state.signals(),
            )
        }
    }

    fn examples(&self) -> Vec<Example> {
        vec![
            Example {
                description: "Extract bytes starting from a specific index",
                example: "{ data: 0x[33 44 55 10 01 13 10] } | bytes at 3.. data",
                result: Some(Value::test_record(record! {
                    "data" => Value::test_binary(vec![0x10, 0x01, 0x13, 0x10]),
                })),
            },
            Example {
                description: "Slice out `0x[10 01 13]` from `0x[33 44 55 10 01 13]`",
                example: "0x[33 44 55 10 01 13] | bytes at 3..5",
                result: Some(Value::test_binary(vec![0x10, 0x01, 0x13])),
            },
            Example {
                description: "Extract bytes from the start up to a specific index",
                example: "0x[33 44 55 10 01 13 10] | bytes at ..4",
                result: Some(Value::test_binary(vec![0x33, 0x44, 0x55, 0x10, 0x01])),
            },
            Example {
                description: "Extract byte `0x[10]` using an exclusive end index",
                example: "0x[33 44 55 10 01 13 10] | bytes at 3..<4",
                result: Some(Value::test_binary(vec![0x10])),
            },
            Example {
                description: "Extract bytes up to a negative index (inclusive)",
                example: "0x[33 44 55 10 01 13 10] | bytes at ..-4",
                result: Some(Value::test_binary(vec![0x33, 0x44, 0x55, 0x10])),
            },
            Example {
                description: "Slice bytes across multiple table columns",
                example: r#"[[ColA ColB ColC]; [0x[11 12 13] 0x[14 15 16] 0x[17 18 19]]] | bytes at 1.. ColB ColC"#,
                result: Some(Value::test_list(vec![Value::test_record(record! {
                    "ColA" => Value::test_binary(vec![0x11, 0x12, 0x13]),
                    "ColB" => Value::test_binary(vec![0x15, 0x16]),
                    "ColC" => Value::test_binary(vec![0x18, 0x19]),
                })])),
            },
            Example {
                description: "Extract the last three bytes using a negative start index",
                example: "0x[33 44 55 10 01 13 10] | bytes at (-3)..",
                result: Some(Value::test_binary(vec![0x01, 0x13, 0x10])),
            },
        ]
    }
}

fn map_value(input: &Value, args: &Arguments, head: Span) -> Value {
    let range = &args.range;
    match input {
        Value::Binary { val, .. } => {
            let len = val.len() as u64;
            let start: u64 = range.absolute_start(len);
            let _start: usize = match start.try_into() {
                Ok(start) => start,
                Err(_) => {
                    let span = input.span();
                    return Value::error(
                        ShellError::UnsupportedInput {
                            msg: format!(
                            "Absolute start position {start} was too large for your system arch."
                        ),
                            input: args.range.to_string(),
                            msg_span: span,
                            input_span: span,
                        },
                        head,
                    );
                }
            };

            let (start, end) = range.absolute_bounds(val.len());
            let bytes: Vec<u8> = match end {
                Bound::Unbounded => val[start..].into(),
                Bound::Included(end) => val[start..=end].into(),
                Bound::Excluded(end) => val[start..end].into(),
            };

            Value::binary(bytes, head)
        }
        Value::Error { .. } => input.clone(),
        other => Value::error(
            ShellError::UnsupportedInput {
                msg: "Only binary values are supported".into(),
                input: format!("input type: {:?}", other.get_type()),
                msg_span: head,
                input_span: other.span(),
            },
            head,
        ),
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_examples() {
        use crate::test_examples;
        test_examples(BytesAt {})
    }
}