diff --git a/crates/nu-command/src/commands/filters/mod.rs b/crates/nu-command/src/commands/filters/mod.rs index 51dfb136ef22..b93404dfaa2f 100644 --- a/crates/nu-command/src/commands/filters/mod.rs +++ b/crates/nu-command/src/commands/filters/mod.rs @@ -38,6 +38,7 @@ mod uniq; mod update; mod where_; mod wrap; +mod zip_; pub use all::Command as All; pub use any::Command as Any; @@ -79,3 +80,4 @@ pub use uniq::Uniq; pub use update::Command as Update; pub use where_::Command as Where; pub use wrap::Wrap; +pub use zip_::Command as Zip; diff --git a/crates/nu-command/src/commands/filters/zip_.rs b/crates/nu-command/src/commands/filters/zip_.rs new file mode 100644 index 000000000000..5757ec15e787 --- /dev/null +++ b/crates/nu-command/src/commands/filters/zip_.rs @@ -0,0 +1,173 @@ +use crate::prelude::*; +use nu_engine::run_block; +use nu_engine::WholeStreamCommand; +use nu_errors::ShellError; +use nu_protocol::did_you_mean; +use nu_protocol::TaggedDictBuilder; +use nu_protocol::{ + hir::CapturedBlock, hir::ExternalRedirection, ColumnPath, PathMember, Signature, SyntaxShape, + UnspannedPathMember, UntaggedValue, Value, +}; +use nu_value_ext::get_data_by_column_path; + +use nu_source::HasFallibleSpan; +pub struct Command; + +impl WholeStreamCommand for Command { + fn name(&self) -> &str { + "zip" + } + + fn signature(&self) -> Signature { + Signature::build("zip").required( + "block", + SyntaxShape::Block, + "the block to run and zip into the table", + ) + } + + fn usage(&self) -> &str { + "Zip two tables." + } + + fn run(&self, args: CommandArgs) -> Result { + command(args) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Zip two lists", + example: "[0 2 4 6 8] | zip { [1 3 5 7 9] } | each { $it }", + result: None, + }, + Example { + description: "Zip two tables", + example: "[[symbol]; ['('] ['['] ['{']] | zip { [[symbol]; [')'] [']'] ['}']] } | each { get symbol | $'($in.0)nushell($in.1)' }", + result: Some(vec![ + Value::from("(nushell)"), + Value::from("[nushell]"), + Value::from("{nushell}") + ]) + }] + } +} + +fn command(args: CommandArgs) -> Result { + let context = &args.context; + let name_tag = args.call_info.name_tag.clone(); + + let block: CapturedBlock = args.req(0)?; + let block_span = &block.block.span.clone(); + let input = args.input; + + context.scope.enter_scope(); + context.scope.add_vars(&block.captured.entries); + let result = run_block( + &block.block, + context, + InputStream::empty(), + ExternalRedirection::Stdout, + ); + context.scope.exit_scope(); + + Ok(OutputStream::from_stream(zip( + input, + result, + name_tag, + *block_span, + )?)) +} + +fn zip<'a>( + l: impl Iterator + 'a + Sync + Send, + r: Result, + command_tag: Tag, + secondary_command_span: Span, +) -> Result + 'a + Sync + Send>, ShellError> { + Ok(Box::new(l.zip(r?).map(move |(s1, s2)| match (s1, s2) { + ( + left_row + @ + Value { + value: UntaggedValue::Row(_), + .. + }, + mut + right_row + @ + Value { + value: UntaggedValue::Row(_), + .. + }, + ) => { + let mut zipped_row = TaggedDictBuilder::new(left_row.tag()); + + right_row.tag = Tag::new(right_row.tag.anchor(), secondary_command_span); + + for column in left_row.data_descriptors() { + let path = ColumnPath::build(&(column.to_string()).spanned(right_row.tag.span)); + zipped_row.insert_value(column, zip_row(&path, &left_row, &right_row)); + } + + zipped_row.into_value() + } + (s1, s2) => { + let mut name_tag = command_tag.clone(); + name_tag.anchor = s1.tag.anchor(); + UntaggedValue::table(&vec![s1, s2]).into_value(&name_tag) + } + }))) +} + +fn zip_row(path: &ColumnPath, left: &Value, right: &Value) -> UntaggedValue { + UntaggedValue::table(&vec![ + get_column(path, left) + .unwrap_or_else(|err| UntaggedValue::Error(err).into_untagged_value()), + get_column(path, right) + .unwrap_or_else(|err| UntaggedValue::Error(err).into_untagged_value()), + ]) +} + +pub fn get_column(path: &ColumnPath, value: &Value) -> Result { + get_data_by_column_path(value, path, move |obj_source, column_path_tried, error| { + let path_members_span = path.maybe_span().unwrap_or_else(Span::unknown); + + if obj_source.is_row() { + if let Some(error) = error_message(column_path_tried, &path_members_span, obj_source) { + return error; + } + } + + error + }) +} + +fn error_message( + column_tried: &PathMember, + path_members_span: &Span, + obj_source: &Value, +) -> Option { + match column_tried { + PathMember { + unspanned: UnspannedPathMember::String(column), + .. + } => { + let primary_label = format!("There isn't a column named '{}' from this table", &column); + + did_you_mean(obj_source, column_tried.as_string()).map(|suggestions| { + ShellError::labeled_error_with_secondary( + "Unknown column", + primary_label, + obj_source.tag.span, + format!( + "Perhaps you meant '{}'? Columns available: {}", + suggestions[0], + &obj_source.data_descriptors().join(", ") + ), + column_tried.span.since(path_members_span), + ) + }) + } + _ => None, + } +} diff --git a/crates/nu-command/src/commands/mod.rs b/crates/nu-command/src/commands/mod.rs index 41c7aa0f115d..755183217809 100644 --- a/crates/nu-command/src/commands/mod.rs +++ b/crates/nu-command/src/commands/mod.rs @@ -108,7 +108,10 @@ mod tests { fn only_examples() -> Vec { let mut commands = full_tests(); - commands.extend(vec![whole_stream_command(Flatten)]); + commands.extend(vec![ + whole_stream_command(Zip), + whole_stream_command(Flatten), + ]); commands } diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index 20c14d21c413..3d54c749a79a 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -184,6 +184,7 @@ pub fn create_default_context(interactive: bool) -> Result Result<(), ShellError> { whole_stream_command(Let {}), whole_stream_command(Select), whole_stream_command(StrCollect), + whole_stream_command(Collect), whole_stream_command(Wrap), cmd, ]); @@ -108,6 +109,7 @@ pub fn test(cmd: impl WholeStreamCommand + 'static) -> Result<(), ShellError> { whole_stream_command(cmd), whole_stream_command(Select), whole_stream_command(StrCollect), + whole_stream_command(Collect), whole_stream_command(Wrap), ]); @@ -182,6 +184,7 @@ pub fn test_dataframe(cmd: impl WholeStreamCommand + 'static) -> Result<(), Shel whole_stream_command(Let {}), whole_stream_command(Select), whole_stream_command(StrCollect), + whole_stream_command(Collect), whole_stream_command(Wrap), whole_stream_command(StrToDatetime), ]); @@ -256,6 +259,7 @@ pub fn test_anchors(cmd: Command) -> Result<(), ShellError> { whole_stream_command(Let {}), whole_stream_command(Select), whole_stream_command(StrCollect), + whole_stream_command(Collect), whole_stream_command(Wrap), cmd, ]); diff --git a/crates/nu-command/tests/commands/mod.rs b/crates/nu-command/tests/commands/mod.rs index c1c7c7718818..4f7a604fa60e 100644 --- a/crates/nu-command/tests/commands/mod.rs +++ b/crates/nu-command/tests/commands/mod.rs @@ -63,3 +63,4 @@ mod where_; mod which; mod with_env; mod wrap; +mod zip; diff --git a/crates/nu-command/tests/commands/zip.rs b/crates/nu-command/tests/commands/zip.rs new file mode 100644 index 000000000000..20aa6bf641e3 --- /dev/null +++ b/crates/nu-command/tests/commands/zip.rs @@ -0,0 +1,77 @@ +use nu_test_support::fs::Stub::FileWithContent; +use nu_test_support::pipeline as input; +use nu_test_support::playground::{says, Playground}; + +use hamcrest2::assert_that; +use hamcrest2::prelude::*; + +const ZIP_POWERED_TEST_ASSERTION_SCRIPT: &str = r#" +def expect [ + left, + right, + --to-eq +] { + $left | zip { $right } | all? { + $it.name.0 == $it.name.1 && $it.commits.0 == $it.commits.1 + } +} + +def add-commits [n] { + each { + let contributor = $it; + let name = $it.name; + let commits = $it.commits; + + $contributor | merge { + [[commits]; [($commits + $n)]] + } + } +} +"#; + +#[test] +fn zips_two_tables() { + Playground::setup("zip_test_1", |dirs, nu| { + nu.with_files(vec![FileWithContent( + "zip_test.nu", + &format!("{}\n", ZIP_POWERED_TEST_ASSERTION_SCRIPT), + )]); + + assert_that!( + nu.pipeline(&input(&format!( + r#" + source {} ; + + let contributors = ([ + [name, commits]; + [andres, 10] + [ jt, 20] + ]); + + let actual = ($contributors | add-commits 10); + + expect $actual --to-eq [[name, commits]; [andres, 20] [jt, 30]] + "#, + dirs.test().join("zip_test.nu").display() + ))), + says().stdout("true") + ); + }) +} + +#[test] +fn zips_two_lists() { + Playground::setup("zip_test_2", |_, nu| { + assert_that!( + nu.pipeline(&input( + r#" + echo [0 2 4 6 8] | zip { [1 3 5 7 9] } + | flatten + | into string + | str collect '-' + "# + )), + says().stdout("0-1-2-3-4-5-6-7-8-9") + ); + }) +}