polars strip-chars: Allow any polars expression for pattern argument (#15178)

# Description 
Allow any polars expression for pattern argument for `polars
strip-chars`
This commit is contained in:
Jack Wright 2025-02-25 15:59:02 -08:00 committed by GitHub
parent 53d30ee7ea
commit 7939fb05ea
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -12,6 +12,7 @@ use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type, Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
Value, Value,
}; };
use polars::prelude::Expr;
#[derive(Clone)] #[derive(Clone)]
pub struct StrStripChars; pub struct StrStripChars;
@ -29,7 +30,11 @@ impl PluginCommand for StrStripChars {
fn signature(&self) -> Signature { fn signature(&self) -> Signature {
Signature::build(self.name()) Signature::build(self.name())
.required("pattern", SyntaxShape::String, "Characters to strip") .required(
"pattern",
SyntaxShape::Any,
"Characters to strip as either a string or polars expression",
)
.switch("start", "Strip from start of strings only", Some('s')) .switch("start", "Strip from start of strings only", Some('s'))
.switch("end", "Strip from end of strings only", Some('e')) .switch("end", "Strip from end of strings only", Some('e'))
.input_output_types(vec![( .input_output_types(vec![(
@ -60,6 +65,26 @@ impl PluginCommand for StrStripChars {
.into_value(Span::test_data()), .into_value(Span::test_data()),
), ),
}, },
Example {
description:
"Strip characters from both ends of strings in a column using an expression",
example: r#"[[text]; ["!!!hello!!!"] ["!!!world!!!"] ["!!!test!!!"]] | polars into-df | polars select (polars col text | polars str-strip-chars (polars lit "!")) | polars collect"#,
result: Some(
NuDataFrame::try_from_columns(
vec![Column::new(
"text".to_string(),
vec![
Value::test_string("hello"),
Value::test_string("world"),
Value::test_string("test"),
],
)],
None,
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
Example { Example {
description: "Strip characters from end of strings in a column", description: "Strip characters from end of strings in a column",
example: r#"[[text]; ["hello!!!"] ["world!!!"] ["test!!!"]] | polars into-df | polars select (polars col text | polars str-strip-chars --end "!") | polars collect"#, example: r#"[[text]; ["hello!!!"] ["world!!!"] ["test!!!"]] | polars into-df | polars select (polars col text | polars str-strip-chars --end "!") | polars collect"#,
@ -123,12 +148,17 @@ fn command_expr(
call: &EvaluatedCall, call: &EvaluatedCall,
expr: NuExpression, expr: NuExpression,
) -> Result<PipelineData, ShellError> { ) -> Result<PipelineData, ShellError> {
let pattern: String = call.req(0)?; let pattern_expr: Expr = call
.req::<Value>(0)
.and_then(|ref v| NuExpression::try_from_value(plugin, v))
.or(call
.req::<String>(0)
.map(polars::prelude::lit)
.map(NuExpression::from))?
.into_polars();
let strip_start = call.has_flag("start")?; let strip_start = call.has_flag("start")?;
let strip_end = call.has_flag("end")?; let strip_end = call.has_flag("end")?;
let pattern_expr = polars::prelude::lit(pattern);
let res: NuExpression = if strip_start { let res: NuExpression = if strip_start {
// Use strip_chars_start when --start flag is provided // Use strip_chars_start when --start flag is provided
expr.into_polars() expr.into_polars()