Mirrors/nushell
2
0
Fork 0
mirror of https://github.com/nushell/nushell synced 2025-01-13 21:55:07 +00:00
Code Issues Projects Releases Packages Wiki Activity

command(split-by) (#897)

Browse source
This commit is contained in:
Andrés N. Robalino 2022-01-30 18:29:21 -05:00 committed by GitHub
parent 76a4455255
commit def5869c1c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 273 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
crates/nu-command/src/default_context.rs
View file

@ -71,6 +71,7 @@ pub fn create_default_context(cwd: impl AsRef<Path>) -> EngineState {
Flatten,
Get,
GroupBy,
SplitBy,
Keep,
Merge,
Move,

2
crates/nu-command/src/filters/mod.rs
View file

@ -33,6 +33,7 @@ mod select;
mod shuffle;
mod skip;
mod sort_by;
mod split_by;
mod transpose;
mod uniq;
mod update;
@ -76,6 +77,7 @@ pub use select::Select;
pub use shuffle::Shuffle;
pub use skip::*;
pub use sort_by::SortBy;
pub use split_by::SplitBy;
pub use transpose::Transpose;
pub use uniq::*;
pub use update::Update;

270
crates/nu-command/src/filters/split_by.rs Normal file
View file

@ -0,0 +1,270 @@
use nu_engine::CallExt;
use nu_protocol::ast::Call;
use nu_protocol::engine::{Command, EngineState, Stack};
use nu_protocol::{
Example, PipelineData, ShellError, Signature, Span, Spanned, SyntaxShape, Value,
};
#[derive(Clone)]
pub struct SplitBy;
impl Command for SplitBy {
fn name(&self) -> &str {
"split-by"
}
fn signature(&self) -> Signature {
Signature::build("split-by").optional(
"splitter",
SyntaxShape::Any,
"the splitter value to use",
)
}
fn usage(&self) -> &str {
"Create a new table splitted."
}
fn run(
&self,
engine_state: &EngineState,
stack: &mut Stack,
call: &Call,
input: PipelineData,
) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
split_by(engine_state, stack, call, input)
}
#[allow(clippy::unwrap_used)]
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "split items by column named \"lang\"",
example: r#"
{
'2019': [
{ name: 'andres', lang: 'rb', year: '2019' },
{ name: 'jt', lang: 'rs', year: '2019' }
],
'2021': [
{ name: 'storm', lang: 'rs', 'year': '2021' }
]
} | split-by lang
"#,
result: Some(Value::Record {
cols: vec!["rb".to_string(), "rs".to_string()],
vals: vec![
Value::Record {
cols: vec!["2019".to_string()],
vals: vec![Value::List {
vals: vec![Value::Record {
cols: vec![
"name".to_string(),
"lang".to_string(),
"year".to_string(),
],
vals: vec![
Value::test_string("andres"),
Value::test_string("rb"),
Value::test_string("2019"),
],
span: Span::test_data(),
}],
span: Span::test_data(),
}],
span: Span::test_data(),
},
Value::Record {
cols: vec!["2019".to_string(), "2021".to_string()],
vals: vec![
Value::List {
vals: vec![Value::Record {
cols: vec![
"name".to_string(),
"lang".to_string(),
"year".to_string(),
],
vals: vec![
Value::test_string("jt"),
Value::test_string("rs"),
Value::test_string("2019"),
],
span: Span::test_data(),
}],
span: Span::test_data(),
},
Value::List {
vals: vec![Value::Record {
cols: vec![
"name".to_string(),
"lang".to_string(),
"year".to_string(),
],
vals: vec![
Value::test_string("storm"),
Value::test_string("rs"),
Value::test_string("2021"),
],
span: Span::test_data(),
}],
span: Span::test_data(),
},
],
span: Span::test_data(),
},
],
span: Span::test_data(),
}),
}]
}
}
enum Grouper {
ByColumn(Option<Spanned<String>>),
}
pub fn split_by(
engine_state: &EngineState,
stack: &mut Stack,
call: &Call,
input: PipelineData,
) -> Result<PipelineData, ShellError> {
let name = call.head;
let splitter: Option<Value> = call.opt(engine_state, stack, 0)?;
match splitter {
Some(v) => {
let splitter = Some(Spanned {
item: v.as_string()?,
span: name,
});
Ok(split(&splitter, input, name)?)
}
None => Err(ShellError::SpannedLabeledError(
"expected name".into(),
"requires a column name for splitting".into(),
name,
)),
}
}
pub fn split(
column_name: &Option<Spanned<String>>,
values: PipelineData,
span: Span,
) -> Result<PipelineData, ShellError> {
let grouper = if let Some(column_name) = column_name {
Grouper::ByColumn(Some(column_name.clone()))
} else {
Grouper::ByColumn(None)
};
match grouper {
Grouper::ByColumn(Some(column_name)) => {
let block =
Box::new(
move |_, row: &Value| match row.get_data_by_key(&column_name.item) {
Some(group_key) => Ok(group_key.as_string()?),
None => Err(ShellError::CantFindColumn(
column_name.span,
row.span().unwrap_or(column_name.span),
)),
},
);
data_split(values, &Some(block), span)
}
Grouper::ByColumn(None) => {
let block = Box::new(move |_, row: &Value| row.as_string());
data_split(values, &Some(block), span)
}
}
}
#[allow(clippy::type_complexity)]
pub fn data_split(
value: PipelineData,
splitter: &Option<Box<dyn Fn(usize, &Value) -> Result<String, ShellError> + Send>>,
span: Span,
) -> Result<PipelineData, ShellError> {
let mut splits = indexmap::IndexMap::new();
let mut cols = vec![];
let mut vals = vec![];
match value {
PipelineData::Value(
Value::Record {
cols,
vals: grouped_rows,
span,
},
_,
) => {
for (idx, list) in grouped_rows.iter().enumerate() {
match super::group_by::data_group(list, splitter, span) {
Ok(grouped) => {
if let Value::Record {
vals: li,
cols: sub_cols,
..
} = grouped
{
for (inner_idx, subset) in li.iter().enumerate() {
let s = splits
.entry(sub_cols[inner_idx].clone())
.or_insert(indexmap::IndexMap::new());
s.insert(cols[idx].clone(), subset.clone());
}
}
}
Err(reason) => return Err(reason),
}
}
}
_ => {
return Err(ShellError::SpannedLabeledError(
"unsupported input".into(),
"requires a table with one row for splitting".into(),
span,
))
}
}
for (k, rows) in splits {
cols.push(k.to_string());
let mut sub_cols = vec![];
let mut sub_vals = vec![];
for (k, v) in rows {
sub_cols.push(k);
sub_vals.push(v);
}
vals.push(Value::Record {
cols: sub_cols,
vals: sub_vals,
span,
});
}
Ok(PipelineData::Value(
Value::Record { cols, vals, span },
None,
))
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_examples() {
use crate::test_examples;
test_examples(SplitBy {})
}
}