Skip to content

Commit 390d06d

Browse files
WindSoilderWindSoilder
authored
add bytes starts-with command (#5950)
* refactor operate, make it generic * refactor operate, add starts with command * add comment * remove useless file
1 parent 89acbda commit 390d06d

4 files changed

Lines changed: 221 additions & 44 deletions

File tree

‎crates/nu-command/src/bytes/length.rs‎

Lines changed: 23 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use super::{operate, BytesArgument};
12
use nu_engine::CallExt;
23
use nu_protocol::ast::Call;
34
use nu_protocol::ast::CellPath;
@@ -8,6 +9,16 @@ use nu_protocol::{Example, PipelineData, ShellError, Signature, Span, SyntaxShap
89
#[derive(Clone)]
910
pub struct BytesLen;
1011

12+
struct Arguments {
13+
column_paths: Option<Vec<CellPath>>,
14+
}
15+
16+
impl BytesArgument for Arguments {
17+
fn take_column_paths(&mut self) -> Option<Vec<CellPath>> {
18+
self.column_paths.take()
19+
}
20+
}
21+
1122
impl Command for BytesLen {
1223
fn name(&self) -> &str {
1324
"bytes length"
@@ -38,7 +49,14 @@ impl Command for BytesLen {
3849
call: &Call,
3950
input: PipelineData,
4051
) -> Result<PipelineData, ShellError> {
41-
operate(engine_state, stack, call, input)
52+
let column_paths: Vec<CellPath> = call.rest(engine_state, stack, 1)?;
53+
let column_paths = if column_paths.is_empty() {
54+
None
55+
} else {
56+
Some(column_paths)
57+
};
58+
let arg = Arguments { column_paths };
59+
operate(length, arg, input, call.head, engine_state.ctrlc.clone())
4260
}
4361

4462
fn examples(&self) -> Vec<Example> {
@@ -60,48 +78,10 @@ impl Command for BytesLen {
6078
}
6179
}
6280

63-
fn operate(
64-
engine_state: &EngineState,
65-
stack: &mut Stack,
66-
call: &Call,
67-
input: PipelineData,
68-
) -> Result<PipelineData, ShellError> {
69-
let head = call.head;
70-
let column_paths: Vec<CellPath> = call.rest(engine_state, stack, 0)?;
71-
if column_paths.is_empty() {
72-
input.map(move |v| action(&v, head), engine_state.ctrlc.clone())
73-
} else {
74-
input.map(
75-
move |mut v| {
76-
for path in &column_paths {
77-
let r =
78-
v.update_cell_path(&path.members, Box::new(move |old| action(old, head)));
79-
if let Err(error) = r {
80-
return Value::Error { error };
81-
}
82-
}
83-
v
84-
},
85-
engine_state.ctrlc.clone(),
86-
)
87-
}
88-
}
89-
90-
fn action(input: &Value, head: Span) -> Value {
91-
match input {
92-
Value::Binary { val, .. } => Value::Int {
93-
val: val.len() as i64,
94-
span: head,
95-
},
96-
other => Value::Error {
97-
error: ShellError::UnsupportedInput(
98-
format!(
99-
"Input's type is {}. This command only works with bytes.",
100-
other.get_type()
101-
),
102-
head,
103-
),
104-
},
81+
fn length(input: &[u8], _arg: &Arguments, span: Span) -> Value {
82+
Value::Int {
83+
val: input.len() as i64,
84+
span,
10585
}
10686
}
10787

‎crates/nu-command/src/bytes/mod.rs‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,77 @@
11
mod length;
2+
mod starts_with;
3+
use nu_protocol::ast::CellPath;
4+
use nu_protocol::{PipelineData, ShellError, Span, Value};
5+
use std::sync::atomic::AtomicBool;
6+
use std::sync::Arc;
27

38
pub use length::BytesLen;
9+
pub use starts_with::BytesStartsWith;
10+
11+
trait BytesArgument {
12+
fn take_column_paths(&mut self) -> Option<Vec<CellPath>>;
13+
}
14+
15+
/// map input pipeline data, for each elements, if it's Binary, invoke relative `cmd` with `arg`.
16+
fn operate<C, A>(
17+
cmd: C,
18+
mut arg: A,
19+
input: PipelineData,
20+
span: Span,
21+
ctrlc: Option<Arc<AtomicBool>>,
22+
) -> Result<PipelineData, ShellError>
23+
where
24+
A: BytesArgument + Send + Sync + 'static,
25+
C: Fn(&[u8], &A, Span) -> Value + Send + Sync + 'static + Clone + Copy,
26+
{
27+
match arg.take_column_paths() {
28+
None => input.map(
29+
move |v| match v {
30+
Value::Binary {
31+
val,
32+
span: val_span,
33+
} => cmd(&val, &arg, val_span),
34+
other => Value::Error {
35+
error: ShellError::UnsupportedInput(
36+
format!(
37+
"Input's type is {}. This command only works with bytes.",
38+
other.get_type()
39+
),
40+
span,
41+
),
42+
},
43+
},
44+
ctrlc,
45+
),
46+
Some(column_paths) => {
47+
let arg = Arc::new(arg);
48+
input.map(
49+
move |mut v| {
50+
for path in &column_paths {
51+
let opt = arg.clone();
52+
let r = v.update_cell_path(
53+
&path.members,
54+
Box::new(move |old| {
55+
match old {
56+
Value::Binary {val, span: val_span} => cmd(val, &opt, *val_span),
57+
other => Value::Error {
58+
error: ShellError::UnsupportedInput(
59+
format!(
60+
"Input's type is {}. This command only works with bytes.",
61+
other.get_type()
62+
),
63+
span,
64+
),
65+
}}}),
66+
);
67+
if let Err(error) = r {
68+
return Value::Error { error };
69+
}
70+
}
71+
v
72+
},
73+
ctrlc,
74+
)
75+
}
76+
}
77+
}

‎crates/nu-command/src/bytes/starts_with.rs‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
use super::{operate, BytesArgument};
2+
use nu_engine::CallExt;
3+
use nu_protocol::ast::Call;
4+
use nu_protocol::ast::CellPath;
5+
use nu_protocol::engine::{Command, EngineState, Stack};
6+
use nu_protocol::Category;
7+
use nu_protocol::{Example, PipelineData, ShellError, Signature, Span, SyntaxShape, Value};
8+
9+
struct Arguments {
10+
pattern: Vec<u8>,
11+
column_paths: Option<Vec<CellPath>>,
12+
}
13+
14+
impl BytesArgument for Arguments {
15+
fn take_column_paths(&mut self) -> Option<Vec<CellPath>> {
16+
self.column_paths.take()
17+
}
18+
}
19+
20+
#[derive(Clone)]
21+
22+
pub struct BytesStartsWith;
23+
24+
impl Command for BytesStartsWith {
25+
fn name(&self) -> &str {
26+
"bytes starts-with"
27+
}
28+
29+
fn signature(&self) -> Signature {
30+
Signature::build("bytes starts-with")
31+
.required("pattern", SyntaxShape::Binary, "the pattern to match")
32+
.rest(
33+
"rest",
34+
SyntaxShape::CellPath,
35+
"optionally matches prefix of text by column paths",
36+
)
37+
.category(Category::Bytes)
38+
}
39+
40+
fn usage(&self) -> &str {
41+
"Check if bytes starts with a pattern"
42+
}
43+
44+
fn search_terms(&self) -> Vec<&str> {
45+
vec!["pattern", "match", "find", "search"]
46+
}
47+
48+
fn run(
49+
&self,
50+
engine_state: &EngineState,
51+
stack: &mut Stack,
52+
call: &Call,
53+
input: PipelineData,
54+
) -> Result<PipelineData, ShellError> {
55+
let pattern: Vec<u8> = call.req(engine_state, stack, 0)?;
56+
let column_paths: Vec<CellPath> = call.rest(engine_state, stack, 1)?;
57+
let column_paths = if column_paths.is_empty() {
58+
None
59+
} else {
60+
Some(column_paths)
61+
};
62+
let arg = Arguments {
63+
pattern,
64+
column_paths,
65+
};
66+
operate(
67+
starts_with,
68+
arg,
69+
input,
70+
call.head,
71+
engine_state.ctrlc.clone(),
72+
)
73+
}
74+
75+
fn examples(&self) -> Vec<Example> {
76+
vec![
77+
Example {
78+
description: "Checks if binary starts with `0x[1F FF AA]`",
79+
example: "0x[1F FF AA AA] | bytes starts-with 0x[1F FF AA]",
80+
result: Some(Value::Bool {
81+
val: true,
82+
span: Span::test_data(),
83+
}),
84+
},
85+
Example {
86+
description: "Checks if binary starts with `0x[1F]`",
87+
example: "0x[1F FF AA AA] | bytes starts-with 0x[1F]",
88+
result: Some(Value::Bool {
89+
val: true,
90+
span: Span::test_data(),
91+
}),
92+
},
93+
Example {
94+
description: "Checks if binary starts with `0x[1F]`",
95+
example: "0x[1F FF AA AA] | bytes starts-with 0x[11]",
96+
result: Some(Value::Bool {
97+
val: false,
98+
span: Span::test_data(),
99+
}),
100+
},
101+
]
102+
}
103+
}
104+
105+
fn starts_with(input: &[u8], Arguments { pattern, .. }: &Arguments, span: Span) -> Value {
106+
Value::Bool {
107+
val: input.starts_with(pattern),
108+
span,
109+
}
110+
}
111+
112+
#[cfg(test)]
113+
mod tests {
114+
use super::*;
115+
116+
#[test]
117+
fn test_examples() {
118+
use crate::test_examples;
119+
120+
test_examples(BytesStartsWith {})
121+
}
122+
}

‎crates/nu-command/src/default_context.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -209,7 +209,8 @@ pub fn create_default_context(cwd: impl AsRef<Path>) -> EngineState {
209209

210210
// Bytes
211211
bind_command! {
212-
BytesLen
212+
BytesLen,
213+
BytesStartsWith
213214
}
214215

215216
// FileSystem

0 commit comments

Comments
 (0)