Port update cells command (#891)

* Port update cells command

Clean up, nicer match statements in UpdateCellsIterator

Return columns flag into HashSet errors

Add FIXME: for update cell behavior on nested lists

* Fix: process cells for Record when no columns are specified

* Fix: address clippy lints for unwrap and into_iter

* Fix: don't step into lists and don't bind $it var
This commit is contained in:
Julian Aichholz 2022-01-30 13:41:05 -08:00 committed by GitHub
parent a51d45b99d
commit 67cb720f24
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 250 additions and 0 deletions

View File

@ -97,6 +97,7 @@ pub fn create_default_context(cwd: impl AsRef<Path>) -> EngineState {
Transpose,
Uniq,
Update,
UpdateCells,
Where,
Wrap,
Zip,

View File

@ -36,6 +36,7 @@ mod sort_by;
mod transpose;
mod uniq;
mod update;
mod update_cells;
mod where_;
mod wrap;
mod zip_;
@ -78,6 +79,7 @@ pub use sort_by::SortBy;
pub use transpose::Transpose;
pub use uniq::*;
pub use update::Update;
pub use update_cells::UpdateCells;
pub use where_::Where;
pub use wrap::Wrap;
pub use zip_::Zip;

View File

@ -0,0 +1,247 @@
use nu_engine::{eval_block, CallExt};
use nu_protocol::ast::{Block, Call};
use nu_protocol::engine::{CaptureBlock, Command, EngineState, Stack};
use nu_protocol::{
Category, Example, IntoInterruptiblePipelineData, IntoPipelineData, PipelineData,
PipelineIterator, ShellError, Signature, Span, SyntaxShape, Value,
};
use std::collections::HashSet;
use std::iter::FromIterator;
#[derive(Clone)]
pub struct UpdateCells;
impl Command for UpdateCells {
fn name(&self) -> &str {
"update cells"
}
fn signature(&self) -> Signature {
Signature::build("update cells")
.required(
"block",
SyntaxShape::Block(Some(vec![SyntaxShape::Any])),
"the block to run an update for each cell",
)
.named(
"columns",
SyntaxShape::Table,
"list of columns to update",
Some('c'),
)
.category(Category::Filters)
}
fn usage(&self) -> &str {
"Update the table cells."
}
fn examples(&self) -> Vec<Example> {
vec![
Example {
description: "Update the zero value cells to empty strings.",
example: r#"[
[2021-04-16, 2021-06-10, 2021-09-18, 2021-10-15, 2021-11-16, 2021-11-17, 2021-11-18];
[ 37, 0, 0, 0, 37, 0, 0]
] | update cells {|value|
if $value == 0 {
""
} else {
$value
}
}"#,
result: Some(Value::List {
vals: vec![Value::Record {
cols: vec![
"2021-04-16".into(),
"2021-06-10".into(),
"2021-09-18".into(),
"2021-10-15".into(),
"2021-11-16".into(),
"2021-11-17".into(),
"2021-11-18".into(),
],
vals: vec![
Value::test_int(37),
Value::test_string(""),
Value::test_string(""),
Value::test_string(""),
Value::test_int(37),
Value::test_string(""),
Value::test_string(""),
],
span: Span::test_data(),
}],
span: Span::test_data(),
}),
},
Example {
description: "Update the zero value cells to empty strings in 2 last columns.",
example: r#"[
[2021-04-16, 2021-06-10, 2021-09-18, 2021-10-15, 2021-11-16, 2021-11-17, 2021-11-18];
[ 37, 0, 0, 0, 37, 0, 0]
] | update cells -c ["2021-11-18", "2021-11-17"] {|value|
if $value == 0 {
""
} else {
$value
}
}"#,
result: Some(Value::List {
vals: vec![Value::Record {
cols: vec![
"2021-04-16".into(),
"2021-06-10".into(),
"2021-09-18".into(),
"2021-10-15".into(),
"2021-11-16".into(),
"2021-11-17".into(),
"2021-11-18".into(),
],
vals: vec![
Value::test_int(37),
Value::test_int(0),
Value::test_int(0),
Value::test_int(0),
Value::test_int(37),
Value::test_string(""),
Value::test_string(""),
],
span: Span::test_data(),
}],
span: Span::test_data(),
}),
},
]
}
fn run(
&self,
engine_state: &EngineState,
stack: &mut Stack,
call: &Call,
input: PipelineData,
) -> Result<PipelineData, ShellError> {
// the block to run on each cell
let engine_state = engine_state.clone();
let block: CaptureBlock = call.req(&engine_state, stack, 0)?;
let mut stack = stack.captures_to_stack(&block.captures);
let orig_env_vars = stack.env_vars.clone();
let orig_env_hidden = stack.env_hidden.clone();
let ctrlc = engine_state.ctrlc.clone();
let block: Block = engine_state.get_block(block.block_id).clone();
let span = call.head;
stack.with_env(&orig_env_vars, &orig_env_hidden);
// the columns to update
let columns: Option<Value> = call.get_flag(&engine_state, &mut stack, "columns")?;
let columns: Option<HashSet<String>> = match columns {
Some(val) => {
let cols = val
.as_list()?
.iter()
.map(|val| val.as_string())
.collect::<Result<Vec<String>, ShellError>>()?;
Some(HashSet::from_iter(cols.into_iter()))
}
None => None,
};
Ok(UpdateCellIterator {
input: input.into_iter(),
engine_state,
stack,
block,
columns,
span,
}
.into_pipeline_data(ctrlc))
}
}
struct UpdateCellIterator {
input: PipelineIterator,
columns: Option<HashSet<String>>,
engine_state: EngineState,
stack: Stack,
block: Block,
span: Span,
}
impl Iterator for UpdateCellIterator {
type Item = Value;
fn next(&mut self) -> Option<Self::Item> {
match self.input.next() {
Some(val) => {
if let Some(ref cols) = self.columns {
if !val.columns().iter().any(|c| cols.contains(c)) {
return Some(val);
}
}
match val {
Value::Record { vals, cols, span } => Some(Value::Record {
vals: cols
.iter()
.zip(vals.into_iter())
.map(|(col, val)| match &self.columns {
Some(cols) if !cols.contains(col) => val,
_ => process_cell(
val,
&self.engine_state,
&mut self.stack,
&self.block,
span,
),
})
.collect(),
cols,
span,
}),
val => Some(process_cell(
val,
&self.engine_state,
&mut self.stack,
&self.block,
self.span,
)),
}
}
None => None,
}
}
}
fn process_cell(
val: Value,
engine_state: &EngineState,
stack: &mut Stack,
block: &Block,
span: Span,
) -> Value {
if let Some(var) = block.signature.get_positional(0) {
if let Some(var_id) = &var.var_id {
stack.add_var(*var_id, val.clone());
}
}
match eval_block(engine_state, stack, block, val.into_pipeline_data()) {
Ok(pd) => pd.into_value(span),
Err(e) => Value::Error { error: e },
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_examples() {
use crate::test_examples;
test_examples(UpdateCells {})
}
}