Port keep, keep while and keep until commands (#384)

* Add `KeepUntil` sub-command

* Add `KeepWhile` sub-command

* Add `Keep` command

* Fix error type

crates/nu-command/src/default_context.rs

@@ -55,6 +55,9 @@ pub fn create_default_context() -> EngineState {
             Each,
             First,
             Get,
+            Keep,
+            KeepUntil,
+            KeepWhile,
             Last,
             Length,
             Lines,

crates/nu-command/src/filters/keep/command.rs

@@ -0,0 +1,100 @@
+use std::convert::TryInto;
+
+use nu_engine::CallExt;
+use nu_protocol::{
+    ast::Call,
+    engine::{Command, EngineState, Stack},
+    Category, Example, IntoInterruptiblePipelineData, PipelineData, ShellError, Signature, Span,
+    SyntaxShape, Value,
+};
+
+#[derive(Clone)]
+pub struct Keep;
+
+impl Command for Keep {
+    fn name(&self) -> &str {
+        "keep"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build(self.name())
+            .optional("n", SyntaxShape::Int, "the number of elements to keep")
+            .category(Category::Filters)
+    }
+
+    fn usage(&self) -> &str {
+        "Keep the first n elements of the input."
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![
+            Example {
+                description: "Keep two elements",
+                example: "echo [[editions]; [2015] [2018] [2021]] | keep 2",
+                result: Some(Value::List {
+                    vals: vec![
+                        Value::Record {
+                            cols: vec!["editions".to_owned()],
+                            vals: vec![Value::from(2015)],
+                            span: Span::unknown(),
+                        },
+                        Value::Record {
+                            cols: vec!["editions".to_owned()],
+                            vals: vec![Value::from(2018)],
+                            span: Span::unknown(),
+                        },
+                    ],
+                    span: Span::unknown(),
+                }),
+            },
+            Example {
+                description: "Keep the first value",
+                example: "echo [2 4 6 8] | keep",
+                result: Some(Value::List {
+                    vals: vec![Value::from(2)],
+                    span: Span::unknown(),
+                }),
+            },
+        ]
+    }
+
+    fn run(
+        &self,
+        engine_state: &EngineState,
+        stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<PipelineData, ShellError> {
+        let n: Option<Value> = call.opt(engine_state, stack, 0)?;
+
+        let n: usize = match n {
+            Some(Value::Int { val, span }) => val.try_into().map_err(|err| {
+                ShellError::UnsupportedInput(
+                    format!("Could not convert {} to unsigned integer: {}", val, err),
+                    span,
+                )
+            })?,
+            Some(_) => {
+                let span = call.head;
+                return Err(ShellError::TypeMismatch("expected integer".into(), span));
+            }
+            None => 1,
+        };
+
+        let ctrlc = engine_state.ctrlc.clone();
+
+        Ok(input.into_iter().take(n).into_pipeline_data(ctrlc))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_examples() {
+        use crate::test_examples;
+
+        test_examples(Keep {})
+    }
+}

crates/nu-command/src/filters/keep/mod.rs

@@ -0,0 +1,7 @@
+mod command;
+mod until;
+mod while_;
+
+pub use command::Keep;
+pub use until::KeepUntil;
+pub use while_::KeepWhile;

crates/nu-command/src/filters/keep/until.rs

@@ -0,0 +1,90 @@
+use nu_engine::eval_block;
+use nu_protocol::{
+    ast::Call,
+    engine::{Command, EngineState, Stack},
+    Category, Example, IntoInterruptiblePipelineData, PipelineData, ShellError, Signature, Span,
+    SyntaxShape, Value,
+};
+
+#[derive(Clone)]
+pub struct KeepUntil;
+
+impl Command for KeepUntil {
+    fn name(&self) -> &str {
+        "keep until"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build(self.name())
+            .required(
+                "predicate",
+                SyntaxShape::RowCondition,
+                "the predicate that kept element must not match",
+            )
+            .category(Category::Filters)
+    }
+
+    fn usage(&self) -> &str {
+        "Keep elements of the input until a predicate is true."
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![Example {
+            description: "Keep until the element is positive",
+            example: "echo [-1 -2 9 1] | keep until $it > 0",
+            result: Some(Value::List {
+                vals: vec![Value::from(-1), Value::from(-2)],
+                span: Span::unknown(),
+            }),
+        }]
+    }
+
+    fn run(
+        &self,
+        engine_state: &EngineState,
+        stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<PipelineData, ShellError> {
+        let span = call.head;
+
+        let predicate = &call.positional[0];
+        let block_id = predicate
+            .as_row_condition_block()
+            .ok_or_else(|| ShellError::TypeMismatch("expected row condition".to_owned(), span))?;
+
+        let block = engine_state.get_block(block_id).clone();
+        let var_id = block.signature.get_positional(0).and_then(|arg| arg.var_id);
+
+        let mut stack = stack.collect_captures(&block.captures);
+
+        let ctrlc = engine_state.ctrlc.clone();
+        let engine_state = engine_state.clone();
+
+        Ok(input
+            .into_iter()
+            .take_while(move |value| {
+                if let Some(var_id) = var_id {
+                    stack.add_var(var_id, value.clone());
+                }
+
+                !eval_block(&engine_state, &mut stack, &block, PipelineData::new(span))
+                    .map_or(false, |pipeline_data| {
+                        pipeline_data.into_value(span).is_true()
+                    })
+            })
+            .into_pipeline_data(ctrlc))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_examples() {
+        use crate::test_examples;
+
+        test_examples(KeepUntil)
+    }
+}

crates/nu-command/src/filters/keep/while_.rs

@@ -0,0 +1,90 @@
+use nu_engine::eval_block;
+use nu_protocol::{
+    ast::Call,
+    engine::{Command, EngineState, Stack},
+    Category, Example, IntoInterruptiblePipelineData, PipelineData, ShellError, Signature, Span,
+    SyntaxShape, Value,
+};
+
+#[derive(Clone)]
+pub struct KeepWhile;
+
+impl Command for KeepWhile {
+    fn name(&self) -> &str {
+        "keep while"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build(self.name())
+            .required(
+                "predicate",
+                SyntaxShape::RowCondition,
+                "the predicate that kept element must not match",
+            )
+            .category(Category::Filters)
+    }
+
+    fn usage(&self) -> &str {
+        "Keep elements of the input while a predicate is true."
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![Example {
+            description: "Keep while the element is negative",
+            example: "echo [-1 -2 9 1] | keep while $it < 0",
+            result: Some(Value::List {
+                vals: vec![Value::from(-1), Value::from(-2)],
+                span: Span::unknown(),
+            }),
+        }]
+    }
+
+    fn run(
+        &self,
+        engine_state: &EngineState,
+        stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<PipelineData, ShellError> {
+        let span = call.head;
+
+        let predicate = &call.positional[0];
+        let block_id = predicate
+            .as_row_condition_block()
+            .ok_or_else(|| ShellError::TypeMismatch("expected row condition".to_owned(), span))?;
+
+        let block = engine_state.get_block(block_id).clone();
+        let var_id = block.signature.get_positional(0).and_then(|arg| arg.var_id);
+
+        let mut stack = stack.collect_captures(&block.captures);
+
+        let ctrlc = engine_state.ctrlc.clone();
+        let engine_state = engine_state.clone();
+
+        Ok(input
+            .into_iter()
+            .take_while(move |value| {
+                if let Some(var_id) = var_id {
+                    stack.add_var(var_id, value.clone());
+                }
+
+                eval_block(&engine_state, &mut stack, &block, PipelineData::new(span))
+                    .map_or(false, |pipeline_data| {
+                        pipeline_data.into_value(span).is_true()
+                    })
+            })
+            .into_pipeline_data(ctrlc))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_examples() {
+        use crate::test_examples;
+
+        test_examples(KeepWhile)
+    }
+}

crates/nu-command/src/filters/mod.rs

@@ -6,6 +6,7 @@ mod drop;
 mod each;
 mod first;
 mod get;
+mod keep;
 mod last;
 mod length;
 mod lines;
@@ -32,6 +33,7 @@ pub use drop::*;
 pub use each::Each;
 pub use first::First;
 pub use get::Get;
+pub use keep::*;
 pub use last::Last;
 pub use length::Length;
 pub use lines::Lines;