use crate::{
PolarsPlugin,
values::{CustomValueSupport, NuDataFrame, NuExpression, PolarsPluginType, str_to_dtype},
};
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::shell_error::generic::GenericError;
use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Value,
};
use polars::{df, prelude::*};
#[derive(Clone)]
pub struct Replace;
impl PluginCommand for Replace {
type Plugin = PolarsPlugin;
fn name(&self) -> &str {
"polars replace"
}
fn description(&self) -> &str {
"Create an expression that replaces old values with new values"
}
fn signature(&self) -> Signature {
Signature::build(self.name())
.required(
"old",
SyntaxShape::OneOf(vec![SyntaxShape::Record(vec![]), SyntaxShape::List(Box::new(SyntaxShape::Any))]),
"Values to be replaced.",
)
.optional(
"new",
SyntaxShape::List(Box::new(SyntaxShape::Any)),
"Values to replace by.",
)
.switch(
"strict",
"Require that all values must be replaced or throw an error (ignored if `old` or `new` are expressions).",
Some('s'),
)
.named(
"default",
SyntaxShape::Any,
"Set values that were not replaced to this value. If no default is specified, (default), an error is raised if any values were not replaced. Accepts expression input. Non-expression inputs are parsed as literals.",
Some('d'),
)
.named(
"return-dtype",
SyntaxShape::String,
"Data type of the resulting expression. If set to `null` (default), the data type is determined automatically based on the other inputs.",
Some('t'),
)
.input_output_type(
PolarsPluginType::NuExpression.into(),
PolarsPluginType::NuExpression.into(),
)
.category(Category::Custom("expression".into()))
}
fn examples(&self) -> Vec<Example<'_>> {
vec![
Example {
description: "Replace column with different values of same type",
example: "[[a]; [1] [1] [2] [2]]
| polars into-df
| polars select (polars col a | polars replace [1 2] [10 20])
| polars collect",
result: Some(
NuDataFrame::from(
df!("a" => [10, 10, 20, 20])
.expect("simple df for test should not fail"),
)
.into_value(Span::test_data()),
),
},
Example {
description: "Replace column with different values of another type",
example: "[[a]; [1] [1] [2] [2]]
| polars into-df
| polars select (polars col a | polars replace [1 2] [a b] --strict)
| polars collect",
result: Some(
NuDataFrame::from(
df!("a" => ["a", "a", "b", "b"])
.expect("simple df for test should not fail"),
)
.into_value(Span::test_data()),
),
},
Example {
description: "Replace column with different values based on expressions (cannot be used with strict)",
example: "[[a]; [1] [1] [2] [2]]
| polars into-df
| polars select (polars col a | polars replace [(polars col a | polars max)] [(polars col a | polars max | $in + 5)])
| polars collect",
result: Some(
NuDataFrame::from(
df!("a" => [1, 1, 7, 7])
.expect("simple df for test should not fail"),
)
.into_value(Span::test_data()),
),
},
Example {
description: "Replace column with different values based on expressions with default",
example: "[[a]; [1] [1] [2] [3]]
| polars into-df
| polars select (polars col a | polars replace [1] [10] --default (polars col a | polars max | $in * 100) --strict)
| polars collect",
result: Some(
NuDataFrame::from(
df!("a" => [10, 10, 300, 300])
.expect("simple df for test should not fail"),
)
.into_value(Span::test_data()),
),
},
Example {
description: "Replace column with different values based on expressions with default",
example: "[[a]; [1] [1] [2] [3]]
| polars into-df
| polars select (polars col a | polars replace [1] [10] --default (polars col a | polars max | $in * 100) --strict --return-dtype str)
| polars collect",
result: Some(
NuDataFrame::from(
df!("a" => ["10", "10", "300", "300"])
.expect("simple df for test should not fail"),
)
.into_value(Span::test_data()),
),
},
Example {
description: "Replace column with different values using a record",
example: "[[a]; [1] [1] [2] [2]]
| polars into-df
| polars select (polars col a | polars replace {1: a, 2: b} --strict --return-dtype str)
| polars collect",
result: Some(
NuDataFrame::from(
df!("a" => ["a", "a", "b", "b"])
.expect("simple df for test should not fail"),
)
.into_value(Span::test_data()),
),
},
]
}
fn search_terms(&self) -> Vec<&str> {
vec!["replace"]
}
fn run(
&self,
plugin: &Self::Plugin,
engine: &EngineInterface,
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let (old_vals, new_vals) = match (call.req(0)?, call.opt::<Value>(1)?) {
(Value::Record { val, .. }, None) => val
.iter()
.map(|(key, value)| (Value::string(key, call.head), value.clone()))
.collect::<Vec<(Value, Value)>>()
.into_iter()
.unzip(),
(Value::List { vals: old_vals, .. }, Some(Value::List { vals: new_vals, .. })) => {
(old_vals, new_vals)
}
(_, _) => {
return Err(LabeledError::from(ShellError::Generic(
GenericError::new("Invalid arguments", "", call.head).with_help("`old` must be either a record or list. If `old` is a record, then `new` must not be specified. Otherwise, `new` must also be a list"),
)));
}
};
let old = values_to_expr(plugin, call.head, old_vals)?;
let new = values_to_expr(plugin, call.head, new_vals)?;
let strict = call.has_flag("strict")?;
let return_dtype = match call.get_flag::<String>("return-dtype")? {
Some(dtype) => {
if !strict {
return Err(LabeledError::from(ShellError::Generic(GenericError::new(
"`return-dtype` may only be used with `strict`",
"",
call.head,
))));
}
Some(str_to_dtype(&dtype, call.head)?)
}
None => None,
};
let default = match call.get_flag::<Value>("default")? {
Some(default) => {
if !strict {
return Err(LabeledError::from(ShellError::Generic(GenericError::new(
"`default` may only be used with `strict`",
"",
call.head,
))));
}
Some(values_to_expr(plugin, call.head, vec![default])?)
}
None => None,
};
let expr = NuExpression::try_from_pipeline(plugin, input, call.head)?;
let expr: NuExpression = if strict {
expr.into_polars()
.replace_strict(old, new, default, return_dtype)
.into()
} else {
expr.into_polars().replace(old, new).into()
};
expr.to_pipeline_data(plugin, engine, call.head)
.map_err(LabeledError::from)
}
}
fn values_to_expr(
plugin: &PolarsPlugin,
span: Span,
values: Vec<Value>,
) -> Result<Expr, ShellError> {
match values.first() {
Some(Value::Int { .. }) => {
let series_values = values
.into_iter()
.filter_map(|v| match v {
Value::Int { val, .. } => Some(val),
_ => None,
})
.collect::<Vec<i64>>();
Ok(lit(Series::new("old".into(), &series_values)))
}
Some(Value::Bool { .. }) => {
let series_values = values
.into_iter()
.filter_map(|v| match v {
Value::Bool { val, .. } => Some(val),
_ => None,
})
.collect::<Vec<bool>>();
Ok(lit(Series::new("old".into(), &series_values)))
}
Some(Value::Float { .. }) => {
let series_values = values
.into_iter()
.filter_map(|v| match v {
Value::Float { val, .. } => Some(val),
_ => None,
})
.collect::<Vec<f64>>();
Ok(lit(Series::new("old".into(), &series_values)))
}
Some(Value::String { .. }) => {
let series_values = values
.into_iter()
.filter_map(|v| match v {
Value::String { val, .. } => Some(val),
_ => None,
})
.collect::<Vec<String>>();
Ok(lit(Series::new("old".into(), &series_values)))
}
Some(Value::Custom { .. }) => {
if values.len() > 1 {
return Err(ShellError::Generic(GenericError::new(
"Multiple expressions to be replaced is not supported",
"",
span,
)));
}
NuExpression::try_from_value(
plugin,
values
.first()
.expect("Presence of first element is enforced at argument parsing."),
)
.map(|expr| expr.into_polars())
}
x @ Some(_) => Err(ShellError::Generic(
GenericError::new("Cannot convert input to expression", "", span)
.with_help(format!("Unexpected type: {x:?}")),
)),
None => Err(ShellError::Generic(GenericError::new(
"Missing input values",
"",
span,
))),
}
}
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command;
#[test]
fn test_examples() -> Result<(), nu_protocol::ShellError> {
test_polars_plugin_command(&Replace)
}
}