use super::super::super::values::{Column, NuDataFrame}; use crate::dataframe::values::NuExpression; use nu_protocol::{ ast::Call, engine::{Command, EngineState, Stack}, Category, Example, PipelineData, ShellError, Signature, Span, Value, }; use polars::prelude::IntoSeries; use std::ops::Not; #[derive(Clone)] pub struct NotSeries; impl Command for NotSeries { fn name(&self) -> &str { "dfr not" } fn usage(&self) -> &str { "Inverts boolean mask or creates a not expression" } fn signature(&self) -> Signature { Signature::build(self.name()).category(Category::Custom("dataframe".into())) } fn examples(&self) -> Vec { vec![ Example { description: "Inverts boolean mask", example: "[true false true] | dfr to-df | dfr not", result: Some( NuDataFrame::try_from_columns(vec![Column::new( "0".to_string(), vec![ Value::test_bool(false), Value::test_bool(true), Value::test_bool(false), ], )]) .expect("simple df for test should not fail") .into_value(Span::test_data()), ), }, Example { description: "Creates a not expression from a column", example: "dfr col a | dfr not", result: None, }, ] } fn run( &self, engine_state: &EngineState, stack: &mut Stack, call: &Call, input: PipelineData, ) -> Result { let value = input.into_value(call.head); if NuExpression::can_downcast(&value) { let expr = NuExpression::try_from_value(value)?; let expr: NuExpression = expr.into_polars().is_null().into(); Ok(PipelineData::Value( NuExpression::into_value(expr, call.head), None, )) } else if NuDataFrame::can_downcast(&value) { let df = NuDataFrame::try_from_value(value)?; command(engine_state, stack, call, df) } else { Err(ShellError::CantConvert( "expression or query".into(), value.get_type().to_string(), value.span()?, None, )) } } } fn command( _engine_state: &EngineState, _stack: &mut Stack, call: &Call, df: NuDataFrame, ) -> Result { let series = df.as_series(call.head)?; let bool = series.bool().map_err(|e| { ShellError::GenericError( "Error inverting mask".into(), e.to_string(), Some(call.head), None, Vec::new(), ) })?; let res = bool.not(); NuDataFrame::try_from_series(vec![res.into_series()], call.head) .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) } #[cfg(test)] mod test { use super::super::super::super::test_dataframe::test_dataframe; use super::*; #[test] fn test_examples() { test_dataframe(vec![Box::new(NotSeries {})]) } }