Skip to content

Commit

Permalink
The zip command. (#3919)
Browse files Browse the repository at this point in the history
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}
───┴───────────
```
  • Loading branch information
andrasio committed Aug 15, 2021
1 parent ee563ec commit b873fa7
Show file tree
Hide file tree
Showing 7 changed files with 264 additions and 3 deletions.
2 changes: 2 additions & 0 deletions crates/nu-command/src/commands/filters/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
173 changes: 173 additions & 0 deletions crates/nu-command/src/commands/filters/zip_.rs
Original file line number Diff line number Diff line change
@@ -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<OutputStream, ShellError> {
command(args)
}

fn examples(&self) -> Vec<Example> {
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<OutputStream, ShellError> {
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<Item = Value> + 'a + Sync + Send,
r: Result<InputStream, ShellError>,
command_tag: Tag,
secondary_command_span: Span,
) -> Result<Box<dyn Iterator<Item = Value> + '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<Value, ShellError> {
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<ShellError> {
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,
}
}
5 changes: 4 additions & 1 deletion crates/nu-command/src/commands/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,10 @@ mod tests {

fn only_examples() -> Vec<Command> {
let mut commands = full_tests();
commands.extend(vec![whole_stream_command(Flatten)]);
commands.extend(vec![
whole_stream_command(Zip),
whole_stream_command(Flatten),
]);
commands
}

Expand Down
1 change: 1 addition & 0 deletions crates/nu-command/src/default_context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,7 @@ pub fn create_default_context(interactive: bool) -> Result<EvaluationContext, Bo
whole_stream_command(RollUp),
whole_stream_command(Rotate),
whole_stream_command(RotateCounterClockwise),
whole_stream_command(Zip),
whole_stream_command(Collect),
// Data processing
whole_stream_command(Histogram),
Expand Down
8 changes: 6 additions & 2 deletions crates/nu-command/src/examples.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,8 @@ use crate::commands::{
};

use crate::commands::{
Append, BuildString, Each, Echo, First, Get, Keep, Last, Let, Math, MathMode, Nth, Select,
StrCollect, Wrap,
Append, BuildString, Collect, Each, Echo, First, Get, Keep, Last, Let, Math, MathMode, Nth,
Select, StrCollect, Wrap,
};
use nu_engine::{run_block, whole_stream_command, Command, EvaluationContext, WholeStreamCommand};
use nu_stream::InputStream;
Expand All @@ -48,6 +48,7 @@ pub fn test_examples(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,
]);
Expand Down Expand Up @@ -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),
]);

Expand Down Expand Up @@ -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),
]);
Expand Down Expand Up @@ -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,
]);
Expand Down
1 change: 1 addition & 0 deletions crates/nu-command/tests/commands/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,3 +63,4 @@ mod where_;
mod which;
mod with_env;
mod wrap;
mod zip;
77 changes: 77 additions & 0 deletions crates/nu-command/tests/commands/zip.rs
Original file line number Diff line number Diff line change
@@ -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")
);
})
}

0 comments on commit b873fa7

Please sign in to comment.