2021-10-02 06:55:05 +02:00
|
|
|
use nu_engine::CallExt;
|
2022-02-09 15:59:40 +01:00
|
|
|
use nu_protocol::ast::{Call, CellPath, PathMember};
|
2021-10-25 18:58:58 +02:00
|
|
|
use nu_protocol::engine::{Command, EngineState, Stack};
|
2021-10-25 06:01:02 +02:00
|
|
|
use nu_protocol::{
|
2022-02-09 15:59:40 +01:00
|
|
|
Category, Example, IntoInterruptiblePipelineData, IntoPipelineData, PipelineData,
|
|
|
|
PipelineIterator, ShellError, Signature, Span, SyntaxShape, Value,
|
2021-10-25 06:01:02 +02:00
|
|
|
};
|
2021-10-02 06:55:05 +02:00
|
|
|
|
2021-10-25 06:01:02 +02:00
|
|
|
#[derive(Clone)]
|
2021-10-02 06:55:05 +02:00
|
|
|
pub struct Select;
|
|
|
|
|
|
|
|
impl Command for Select {
|
|
|
|
fn name(&self) -> &str {
|
|
|
|
"select"
|
|
|
|
}
|
|
|
|
|
2022-02-09 15:59:40 +01:00
|
|
|
// FIXME: also add support for --skip
|
2021-10-02 06:55:05 +02:00
|
|
|
fn signature(&self) -> Signature {
|
2021-11-17 05:22:37 +01:00
|
|
|
Signature::build("select")
|
|
|
|
.rest(
|
|
|
|
"rest",
|
|
|
|
SyntaxShape::CellPath,
|
|
|
|
"the columns to select from the table",
|
|
|
|
)
|
|
|
|
.category(Category::Filters)
|
2021-10-02 06:55:05 +02:00
|
|
|
}
|
|
|
|
|
|
|
|
fn usage(&self) -> &str {
|
|
|
|
"Down-select table to only these columns."
|
|
|
|
}
|
|
|
|
|
|
|
|
fn run(
|
|
|
|
&self,
|
2021-10-25 08:31:39 +02:00
|
|
|
engine_state: &EngineState,
|
|
|
|
stack: &mut Stack,
|
2021-10-02 06:55:05 +02:00
|
|
|
call: &Call,
|
2021-10-25 06:01:02 +02:00
|
|
|
input: PipelineData,
|
|
|
|
) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
|
2021-10-25 08:31:39 +02:00
|
|
|
let columns: Vec<CellPath> = call.rest(engine_state, stack, 0)?;
|
2021-10-02 06:55:05 +02:00
|
|
|
let span = call.head;
|
|
|
|
|
2021-10-28 06:13:10 +02:00
|
|
|
select(engine_state, span, columns, input)
|
2021-10-02 06:55:05 +02:00
|
|
|
}
|
|
|
|
|
|
|
|
fn examples(&self) -> Vec<Example> {
|
|
|
|
vec![
|
|
|
|
Example {
|
|
|
|
description: "Select just the name column",
|
|
|
|
example: "ls | select name",
|
|
|
|
result: None,
|
|
|
|
},
|
|
|
|
Example {
|
|
|
|
description: "Select the name and size columns",
|
|
|
|
example: "ls | select name size",
|
|
|
|
result: None,
|
|
|
|
},
|
|
|
|
]
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2021-10-25 06:01:02 +02:00
|
|
|
fn select(
|
2021-10-28 06:13:10 +02:00
|
|
|
engine_state: &EngineState,
|
2021-10-25 06:01:02 +02:00
|
|
|
span: Span,
|
|
|
|
columns: Vec<CellPath>,
|
|
|
|
input: PipelineData,
|
|
|
|
) -> Result<PipelineData, ShellError> {
|
2022-02-09 15:59:40 +01:00
|
|
|
let mut rows = vec![];
|
|
|
|
|
|
|
|
let mut new_columns = vec![];
|
|
|
|
|
|
|
|
for column in columns {
|
|
|
|
let CellPath { ref members } = column;
|
|
|
|
match members.get(0) {
|
|
|
|
Some(PathMember::Int { val, span }) => {
|
|
|
|
if members.len() > 1 {
|
|
|
|
return Err(ShellError::SpannedLabeledError(
|
|
|
|
"Select only allows row numbers for rows".into(),
|
|
|
|
"extra after row number".into(),
|
|
|
|
*span,
|
|
|
|
));
|
|
|
|
}
|
|
|
|
|
|
|
|
rows.push(*val);
|
|
|
|
}
|
|
|
|
_ => new_columns.push(column),
|
|
|
|
};
|
2021-10-02 06:55:05 +02:00
|
|
|
}
|
2022-02-09 15:59:40 +01:00
|
|
|
let columns = new_columns;
|
|
|
|
|
|
|
|
let input = if !rows.is_empty() {
|
|
|
|
rows.sort_unstable();
|
|
|
|
// let skip = call.has_flag("skip");
|
|
|
|
let pipeline_iter: PipelineIterator = input.into_iter();
|
|
|
|
|
|
|
|
NthIterator {
|
|
|
|
input: pipeline_iter,
|
|
|
|
rows,
|
|
|
|
skip: false,
|
|
|
|
current: 0,
|
|
|
|
}
|
|
|
|
.into_pipeline_data(engine_state.ctrlc.clone())
|
|
|
|
} else {
|
|
|
|
input
|
|
|
|
};
|
2021-10-02 06:55:05 +02:00
|
|
|
|
|
|
|
match input {
|
2021-12-02 06:59:10 +01:00
|
|
|
PipelineData::Value(
|
|
|
|
Value::List {
|
|
|
|
vals: input_vals,
|
|
|
|
span,
|
|
|
|
},
|
|
|
|
..,
|
|
|
|
) => {
|
2021-10-02 06:55:05 +02:00
|
|
|
let mut output = vec![];
|
|
|
|
|
|
|
|
for input_val in input_vals {
|
2022-02-09 15:59:40 +01:00
|
|
|
if !columns.is_empty() {
|
|
|
|
let mut cols = vec![];
|
|
|
|
let mut vals = vec![];
|
|
|
|
for path in &columns {
|
|
|
|
//FIXME: improve implementation to not clone
|
|
|
|
let fetcher = input_val.clone().follow_cell_path(&path.members)?;
|
|
|
|
|
|
|
|
cols.push(path.into_string());
|
|
|
|
vals.push(fetcher);
|
|
|
|
}
|
2021-10-02 06:55:05 +02:00
|
|
|
|
2022-02-09 15:59:40 +01:00
|
|
|
output.push(Value::Record { cols, vals, span })
|
|
|
|
} else {
|
|
|
|
output.push(input_val)
|
2021-10-02 06:55:05 +02:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2021-10-28 06:13:10 +02:00
|
|
|
Ok(output
|
|
|
|
.into_iter()
|
|
|
|
.into_pipeline_data(engine_state.ctrlc.clone()))
|
2021-10-02 06:55:05 +02:00
|
|
|
}
|
2021-12-24 08:22:11 +01:00
|
|
|
PipelineData::ListStream(stream, ..) => Ok(stream
|
2021-10-25 06:01:02 +02:00
|
|
|
.map(move |x| {
|
2022-02-09 15:59:40 +01:00
|
|
|
if !columns.is_empty() {
|
|
|
|
let mut cols = vec![];
|
|
|
|
let mut vals = vec![];
|
|
|
|
for path in &columns {
|
|
|
|
//FIXME: improve implementation to not clone
|
|
|
|
match x.clone().follow_cell_path(&path.members) {
|
|
|
|
Ok(value) => {
|
|
|
|
cols.push(path.into_string());
|
|
|
|
vals.push(value);
|
|
|
|
}
|
2022-03-09 13:05:55 +01:00
|
|
|
Err(_) => {
|
2022-02-09 15:59:40 +01:00
|
|
|
cols.push(path.into_string());
|
2022-03-09 13:05:55 +01:00
|
|
|
vals.push(Value::Nothing { span });
|
2022-02-09 15:59:40 +01:00
|
|
|
}
|
2021-10-02 06:55:05 +02:00
|
|
|
}
|
|
|
|
}
|
2022-02-09 15:59:40 +01:00
|
|
|
Value::Record { cols, vals, span }
|
|
|
|
} else {
|
|
|
|
x
|
2021-10-25 06:01:02 +02:00
|
|
|
}
|
|
|
|
})
|
2021-10-28 06:13:10 +02:00
|
|
|
.into_pipeline_data(engine_state.ctrlc.clone())),
|
2021-12-02 06:59:10 +01:00
|
|
|
PipelineData::Value(v, ..) => {
|
2022-02-09 15:59:40 +01:00
|
|
|
if !columns.is_empty() {
|
|
|
|
let mut cols = vec![];
|
|
|
|
let mut vals = vec![];
|
2021-10-02 06:55:05 +02:00
|
|
|
|
2022-02-09 15:59:40 +01:00
|
|
|
for cell_path in columns {
|
|
|
|
// FIXME: remove clone
|
|
|
|
let result = v.clone().follow_cell_path(&cell_path.members)?;
|
2021-10-02 06:55:05 +02:00
|
|
|
|
2022-02-09 15:59:40 +01:00
|
|
|
cols.push(cell_path.into_string());
|
|
|
|
vals.push(result);
|
|
|
|
}
|
2021-10-02 06:55:05 +02:00
|
|
|
|
2022-02-09 15:59:40 +01:00
|
|
|
Ok(Value::Record { cols, vals, span }.into_pipeline_data())
|
|
|
|
} else {
|
|
|
|
Ok(v.into_pipeline_data())
|
|
|
|
}
|
2021-10-02 06:55:05 +02:00
|
|
|
}
|
2021-12-24 08:22:11 +01:00
|
|
|
_ => Ok(PipelineData::new(span)),
|
2021-10-02 06:55:05 +02:00
|
|
|
}
|
|
|
|
}
|
2022-02-09 15:59:40 +01:00
|
|
|
|
|
|
|
struct NthIterator {
|
|
|
|
input: PipelineIterator,
|
|
|
|
rows: Vec<usize>,
|
|
|
|
skip: bool,
|
|
|
|
current: usize,
|
|
|
|
}
|
|
|
|
|
|
|
|
impl Iterator for NthIterator {
|
|
|
|
type Item = Value;
|
|
|
|
|
|
|
|
fn next(&mut self) -> Option<Self::Item> {
|
|
|
|
loop {
|
|
|
|
if !self.skip {
|
|
|
|
if let Some(row) = self.rows.get(0) {
|
|
|
|
if self.current == *row {
|
|
|
|
self.rows.remove(0);
|
|
|
|
self.current += 1;
|
|
|
|
return self.input.next();
|
|
|
|
} else {
|
|
|
|
self.current += 1;
|
|
|
|
let _ = self.input.next();
|
|
|
|
continue;
|
|
|
|
}
|
|
|
|
} else {
|
|
|
|
return None;
|
|
|
|
}
|
|
|
|
} else if let Some(row) = self.rows.get(0) {
|
|
|
|
if self.current == *row {
|
|
|
|
self.rows.remove(0);
|
|
|
|
self.current += 1;
|
|
|
|
let _ = self.input.next();
|
|
|
|
continue;
|
|
|
|
} else {
|
|
|
|
self.current += 1;
|
|
|
|
return self.input.next();
|
|
|
|
}
|
|
|
|
} else {
|
|
|
|
return self.input.next();
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|