diff --git a/crates/nu-command/src/strings/split/list.rs b/crates/nu-command/src/strings/split/list.rs index d960b9a515..cf7cd5ef2a 100644 --- a/crates/nu-command/src/strings/split/list.rs +++ b/crates/nu-command/src/strings/split/list.rs @@ -1,5 +1,6 @@ use fancy_regex::Regex; -use nu_engine::command_prelude::*; +use nu_engine::{command_prelude::*, ClosureEval}; +use nu_protocol::{FromValue, Signals}; #[derive(Clone)] pub struct SubCommand; @@ -24,6 +25,7 @@ impl Command for SubCommand { "regex", "separator is a regular expression, matching values that can be coerced into a string", Some('r')) + .named("split", SyntaxShape::String, "Whether to split lists before, after, or on (default) the separator", None) .category(Category::Filters) } @@ -90,6 +92,7 @@ impl Command for SubCommand { example: "[a, b, c, d, a, e, f, g] | split list a", result: Some(Value::list( vec![ + Value::list(vec![], Span::test_data()), Value::list( vec![ Value::test_string("b"), @@ -132,6 +135,36 @@ impl Command for SubCommand { Span::test_data(), )), }, + Example { + description: "Split a list of numbers on multiples of 3", + example: r"[1 2 3 4 5 6 7 8 9 10] | split list {|e| $e mod 3 == 0 }", + result: Some(Value::test_list(vec![ + Value::test_list(vec![Value::test_int(1), Value::test_int(2)]), + Value::test_list(vec![Value::test_int(4), Value::test_int(5)]), + Value::test_list(vec![Value::test_int(7), Value::test_int(8)]), + Value::test_list(vec![Value::test_int(10)]), + ])), + }, + Example { + description: "Split a list of numbers into lists ending with 0", + example: r"[1 2 0 3 4 5 0 6 0 0 7] | split list --split after 0", + result: Some(Value::test_list(vec![ + Value::test_list(vec![ + Value::test_int(1), + Value::test_int(2), + Value::test_int(0), + ]), + Value::test_list(vec![ + Value::test_int(3), + Value::test_int(4), + Value::test_int(5), + Value::test_int(0), + ]), + Value::test_list(vec![Value::test_int(6), Value::test_int(0)]), + Value::test_list(vec![Value::test_int(0)]), + Value::test_list(vec![Value::test_int(7)]), + ])), + }, ] } @@ -148,7 +181,15 @@ impl Command for SubCommand { ) -> Result { let has_regex = call.has_flag(engine_state, stack, "regex")?; let separator: Value = call.req(engine_state, stack, 0)?; - split_list(engine_state, call, input, has_regex, separator) + let split: Option = call.get_flag(engine_state, stack, "split")?; + let split = split.unwrap_or(Split::On); + let matcher = match separator { + Value::Closure { val, .. } => { + Matcher::from_closure(ClosureEval::new(engine_state, stack, *val)) + } + _ => Matcher::new(has_regex, separator)?, + }; + split_list(engine_state, call, input, matcher, split) } fn run_const( @@ -159,13 +200,40 @@ impl Command for SubCommand { ) -> Result { let has_regex = call.has_flag_const(working_set, "regex")?; let separator: Value = call.req_const(working_set, 0)?; - split_list(working_set.permanent(), call, input, has_regex, separator) + let split: Option = call.get_flag_const(working_set, "split")?; + let split = split.unwrap_or(Split::On); + let matcher = Matcher::new(has_regex, separator)?; + split_list(working_set.permanent(), call, input, matcher, split) } } enum Matcher { Regex(Regex), Direct(Value), + Closure(ClosureEval), +} + +enum Split { + On, + Before, + After, +} + +impl FromValue for Split { + fn from_value(v: Value) -> Result { + let span = v.span(); + let s = ::from_value(v)?; + match s.as_str() { + "on" => Ok(Split::On), + "before" => Ok(Split::Before), + "after" => Ok(Split::After), + _ => Err(ShellError::InvalidValue { + valid: "one of: on, before, after".into(), + actual: s, + span, + }), + } + } } impl Matcher { @@ -188,7 +256,11 @@ impl Matcher { } } - pub fn compare(&self, rhs: &Value) -> Result { + pub fn from_closure(closure: ClosureEval) -> Self { + Self::Closure(closure) + } + + pub fn compare(&mut self, rhs: &Value) -> Result { Ok(match self { Matcher::Regex(regex) => { if let Ok(rhs_str) = rhs.coerce_str() { @@ -198,6 +270,11 @@ impl Matcher { } } Matcher::Direct(lhs) => rhs == lhs, + Matcher::Closure(closure) => closure + .run_with_value(rhs.clone()) + .and_then(|data| data.into_value(Span::unknown())) + .map(|value| value.is_true()) + .unwrap_or(false), }) } } @@ -206,29 +283,96 @@ fn split_list( engine_state: &EngineState, call: &Call, input: PipelineData, - has_regex: bool, - separator: Value, + mut matcher: Matcher, + split: Split, ) -> Result { - let mut temp_list = Vec::new(); - let mut returned_list = Vec::new(); + let head = call.head; + Ok(SplitList::new( + input.into_iter(), + engine_state.signals().clone(), + split, + move |x| matcher.compare(x).unwrap_or(false), + ) + .map(move |x| Value::list(x, head)) + .into_pipeline_data(head, engine_state.signals().clone())) +} - let matcher = Matcher::new(has_regex, separator)?; - for val in input { - engine_state.signals().check(call.head)?; +struct SplitList { + iterator: I, + closure: F, + done: bool, + signals: Signals, + split: Split, + last_item: Option, +} - if matcher.compare(&val)? { - if !temp_list.is_empty() { - returned_list.push(Value::list(temp_list.clone(), call.head)); - temp_list = Vec::new(); - } - } else { - temp_list.push(val); +impl SplitList +where + I: Iterator, + F: FnMut(&I::Item) -> bool, +{ + fn new(iterator: I, signals: Signals, split: Split, closure: F) -> Self { + Self { + iterator, + closure, + done: false, + signals, + split, + last_item: None, } } - if !temp_list.is_empty() { - returned_list.push(Value::list(temp_list.clone(), call.head)); + + fn inner_iterator_next(&mut self) -> Option { + if self.signals.interrupted() { + self.done = true; + return None; + } + self.iterator.next() + } +} + +impl Iterator for SplitList +where + I: Iterator, + F: FnMut(&I::Item) -> bool, +{ + type Item = Vec; + + fn next(&mut self) -> Option { + if self.done { + return None; + } + + let mut items = vec![]; + if let Some(item) = self.last_item.take() { + items.push(item); + } + + loop { + match self.inner_iterator_next() { + None => { + self.done = true; + return Some(items); + } + Some(value) => { + if (self.closure)(&value) { + match self.split { + Split::On => {} + Split::Before => { + self.last_item = Some(value); + } + Split::After => { + items.push(value); + } + } + return Some(items); + } else { + items.push(value); + } + } + } + } } - Ok(Value::list(returned_list, call.head).into_pipeline_data()) } #[cfg(test)]