From b873fa7a5f6e46986520a367e2c8d86c22abc9a8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andr=C3=A9s=20N=2E=20Robalino?= Date: Sat, 14 Aug 2021 23:36:08 -0500 Subject: [PATCH] The zip command. (#3919) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit We introduce it here and allow it to work with regular lists (tables with no columns) as well as symmetric tables. Say we have two lists and wish to zip them, like so: ``` [0 2 4 6 8] | zip { [1 3 5 7 9] } | flatten ───┬─── 0 │ 0 1 │ 1 2 │ 2 3 │ 3 4 │ 4 5 │ 5 6 │ 6 7 │ 7 8 │ 8 9 │ 9 ───┴─── ``` In the case for two tables instead: ``` [[symbol]; ['('] ['['] ['{']] | zip { [[symbol]; [')'] [']'] ['}']] } | each { get symbol | $'($in.0)nushell($in.1)' } ───┬─────────── 0 │ (nushell) 1 │ [nushell] 2 │ {nushell} ───┴─────────── ``` --- crates/nu-command/src/commands/filters/mod.rs | 2 + .../nu-command/src/commands/filters/zip_.rs | 173 ++++++++++++++++++ crates/nu-command/src/commands/mod.rs | 5 +- crates/nu-command/src/default_context.rs | 1 + crates/nu-command/src/examples.rs | 8 +- crates/nu-command/tests/commands/mod.rs | 1 + crates/nu-command/tests/commands/zip.rs | 77 ++++++++ 7 files changed, 264 insertions(+), 3 deletions(-) create mode 100644 crates/nu-command/src/commands/filters/zip_.rs create mode 100644 crates/nu-command/tests/commands/zip.rs 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") + ); + }) +}