use crate::{
PolarsPlugin,
dataframe::values::NuSelector,
values::{CustomValueSupport, PolarsPluginType},
};
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{Category, Example, LabeledError, PipelineData, Signature, Type};
use polars::prelude::{DataTypeSelector, Selector};
#[derive(Clone)]
pub struct SelectorNested;
impl PluginCommand for SelectorNested {
type Plugin = PolarsPlugin;
fn name(&self) -> &str {
"polars selector nested"
}
fn description(&self) -> &str {
"Select all nested columns (list, array, or struct)."
}
fn signature(&self) -> Signature {
Signature::build(self.name())
.input_output_type(Type::Any, PolarsPluginType::NuSelector.into())
.category(Category::Custom("expression".into()))
}
fn examples(&self) -> Vec<Example<'_>> {
vec![Example {
example: "polars selector nested",
description: "Create a selector for nested columns (list, array, or struct)",
result: None,
}]
}
fn search_terms(&self) -> Vec<&str> {
vec!["columns", "select", "nested", "list", "array", "struct"]
}
fn run(
&self,
plugin: &Self::Plugin,
engine: &EngineInterface,
call: &EvaluatedCall,
mut input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let metadata = input.take_metadata();
let selector = Selector::ByDType(DataTypeSelector::Nested);
let nu_selector = NuSelector::from(selector);
nu_selector
.to_pipeline_data(plugin, engine, call.head)
.map_err(LabeledError::from)
.map(|pd| pd.set_metadata(metadata))
}
}
#[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(&SelectorNested)
}
}