Add split number flag in split row (#5434)

Signed-off-by: Yuheng Su <gipsyh.icu@gmail.com>
This commit is contained in:
Yuheng Su 2022-05-06 23:53:02 +08:00 committed by GitHub
parent c2ea993f7e
commit fbdb125141
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -20,6 +20,12 @@ impl Command for SubCommand {
SyntaxShape::String, SyntaxShape::String,
"the character that denotes what separates rows", "the character that denotes what separates rows",
) )
.named(
"number",
SyntaxShape::Int,
"Split into maximum number of items",
Some('n'),
)
.category(Category::Strings) .category(Category::Strings)
} }
@ -75,26 +81,44 @@ fn split_row(
) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> { ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
let name_span = call.head; let name_span = call.head;
let separator: Spanned<String> = call.req(engine_state, stack, 0)?; let separator: Spanned<String> = call.req(engine_state, stack, 0)?;
let max_split: Option<usize> = call.get_flag(engine_state, stack, "number")?;
input.flat_map( input.flat_map(
move |x| split_row_helper(&x, &separator, name_span), move |x| split_row_helper(&x, &separator, max_split, name_span),
engine_state.ctrlc.clone(), engine_state.ctrlc.clone(),
) )
} }
fn split_row_helper(v: &Value, separator: &Spanned<String>, name: Span) -> Vec<Value> { fn split_row_helper(
v: &Value,
separator: &Spanned<String>,
max_split: Option<usize>,
name: Span,
) -> Vec<Value> {
match v.span() { match v.span() {
Ok(v_span) => { Ok(v_span) => {
if let Ok(s) = v.as_string() { if let Ok(s) = v.as_string() {
s.split(&separator.item) match max_split {
.filter_map(|s| { Some(max_split) => s
if s.trim() != "" { .splitn(max_split, &separator.item)
Some(Value::string(s, v_span)) .filter_map(|s| {
} else { if s.trim() != "" {
None Some(Value::string(s, v_span))
} } else {
}) None
.collect() }
})
.collect(),
None => s
.split(&separator.item)
.filter_map(|s| {
if s.trim() != "" {
Some(Value::string(s, v_span))
} else {
None
}
})
.collect(),
}
} else { } else {
vec![Value::Error { vec![Value::Error {
error: ShellError::PipelineMismatch("string".into(), name, v_span), error: ShellError::PipelineMismatch("string".into(), name, v_span),