diff --git a/crates/nu-command/src/filters/reject.rs b/crates/nu-command/src/filters/reject.rs index fdb63fcfe1..a790806099 100644 --- a/crates/nu-command/src/filters/reject.rs +++ b/crates/nu-command/src/filters/reject.rs @@ -24,7 +24,10 @@ impl Command for Reject { ]) .rest( "rest", - SyntaxShape::CellPath, + SyntaxShape::OneOf(vec![ + SyntaxShape::CellPath, + SyntaxShape::List(Box::new(SyntaxShape::CellPath)), + ]), "the names of columns to remove from the table", ) .category(Category::Filters) @@ -49,9 +52,81 @@ impl Command for Reject { call: &Call, input: PipelineData, ) -> Result { - let columns: Vec = call.rest(engine_state, stack, 0)?; + let columns: Vec = call.rest(engine_state, stack, 0)?; + let mut new_columns: Vec = vec![]; + for col_val in columns { + let col_span = &col_val.span(); + match col_val { + Value::CellPath { val, .. } => { + new_columns.push(val); + } + Value::List { vals, .. } => { + for value in vals { + let val_span = &value.span(); + match value { + Value::String { val, .. } => { + let cv = CellPath { + members: vec![PathMember::String { + val: val.clone(), + span: *val_span, + optional: false, + }], + }; + new_columns.push(cv.clone()); + } + Value::Int { val, .. } => { + let cv = CellPath { + members: vec![PathMember::Int { + val: val as usize, + span: *val_span, + optional: false, + }], + }; + new_columns.push(cv.clone()); + } + y => { + return Err(ShellError::CantConvert { + to_type: "cell path".into(), + from_type: y.get_type().to_string(), + span: y.span(), + help: None, + }); + } + } + } + } + Value::String { val, .. } => { + let cv = CellPath { + members: vec![PathMember::String { + val: val.clone(), + span: *col_span, + optional: false, + }], + }; + new_columns.push(cv.clone()); + } + Value::Int { val, .. } => { + let cv = CellPath { + members: vec![PathMember::Int { + val: val as usize, + span: *col_span, + optional: false, + }], + }; + new_columns.push(cv.clone()); + } + x => { + return Err(ShellError::CantConvert { + to_type: "cell path".into(), + from_type: x.get_type().to_string(), + span: x.span(), + help: None, + }); + } + } + } let span = call.head; - reject(engine_state, span, input, columns) + reject(engine_state, span, input, new_columns) } fn examples(&self) -> Vec { @@ -75,13 +150,13 @@ impl Command for Reject { Example { description: "Reject a row in a table", example: "[[a, b]; [1, 2] [3, 4]] | reject 1", - result: Some(Value::List { - vals: vec![Value::test_record(Record { + result: Some(Value::list( + vec![Value::test_record(Record { cols: vec!["a".to_string(), "b".to_string()], vals: vec![Value::test_int(1), Value::test_int(2)], })], - span: Span::test_data(), - }), + Span::test_data(), + )), }, Example { description: "Reject the specified field in a record", @@ -102,6 +177,16 @@ impl Command for Reject { })], })), }, + Example { + description: "Reject columns by a provided list of columns", + example: "let cols = [size type];[[name type size]; [Cargo.toml toml 1kb] [Cargo.lock toml 2kb]] | reject $cols", + result: None + }, + Example { + description: "Reject rows by a provided list of rows", + example: "let rows = [0 2];[[name type size]; [Cargo.toml toml 1kb] [Cargo.lock toml 2kb] [file.json json 3kb]] | reject $rows", + result: None + }, ] } } diff --git a/crates/nu-command/tests/commands/reject.rs b/crates/nu-command/tests/commands/reject.rs index d073a9be30..ecefc15dce 100644 --- a/crates/nu-command/tests/commands/reject.rs +++ b/crates/nu-command/tests/commands/reject.rs @@ -150,6 +150,30 @@ fn reject_optional_row() { assert_eq!(actual.out, "[[foo]; [bar]]"); } +#[test] +fn reject_list_columns() { + let actual = nu!("let arg = [type size]; [[name type size];[Cargo.toml file 10mb] [Cargo.lock file 10mb] [src dir 100mb]] | reject $arg | to nuon"); + assert_eq!(actual.out, "[[name]; [Cargo.toml], [Cargo.lock], [src]]"); +} + +#[test] +fn reject_list_rows() { + let actual = nu!("let arg = [2 0]; [[name type size];[Cargo.toml file 10mb] [Cargo.lock file 10mb] [src dir 100mb]] | reject $arg | to nuon"); + assert_eq!( + actual.out, + "[[name, type, size]; [Cargo.lock, file, 10000000b]]" + ); +} + +#[test] +fn rject_list_mixed() { + let actual = nu!("let arg = [ type 2]; [[name type size];[Cargp.toml file 10mb] [ Cargo.lock file 10mb] [src dir 100mb]] | reject $arg | to nuon"); + assert_eq!( + actual.out, + "[[name, size]; [Cargp.toml, 10000000b], [Cargo.lock, 10000000b]]" + ); +} + #[test] fn reject_multiple_rows_ascending() { let actual = nu!("[[a,b];[1 2] [3 4] [5 6]] | reject 1 2 | to nuon");