use crate::values::{Column, NuDataFrame, NuSchema};
use crate::{
PolarsPlugin,
dataframe::values::NuExpression,
values::{CustomValueSupport, PolarsPluginObject, PolarsPluginType, cant_convert_err},
};
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::shell_error::generic::GenericError;
use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Spanned,
SyntaxShape, Value,
};
use chrono::DateTime;
use polars::prelude::*;
use polars_plan::plans::DynLiteralValue;
use super::timezone_from_str;
#[derive(Clone)]
pub struct ReplaceTimeZone;
impl PluginCommand for ReplaceTimeZone {
type Plugin = PolarsPlugin;
fn name(&self) -> &str {
"polars replace-time-zone"
}
fn description(&self) -> &str {
"Replace the timezone information in a datetime column."
}
fn signature(&self) -> Signature {
Signature::build(self.name())
.input_output_types(vec![
(
PolarsPluginType::NuExpression.into(),
PolarsPluginType::NuExpression.into(),
),
(
PolarsPluginType::NuSelector.into(),
PolarsPluginType::NuExpression.into(),
),
])
.required(
"time_zone",
SyntaxShape::String,
"Timezone for the Datetime Series. Pass `null` to unset time zone.",
)
.named(
"ambiguous",
SyntaxShape::OneOf(vec![SyntaxShape::String, SyntaxShape::Nothing]),
r#"Determine how to deal with ambiguous datetimes:
`raise` (default): raise error
`earliest`: use the earliest datetime
`latest`: use the latest datetime
`null`: set to null"#,
Some('a'),
)
.named(
"nonexistent",
SyntaxShape::OneOf(vec![SyntaxShape::String, SyntaxShape::Nothing]),
r#"Determine how to deal with non-existent datetimes: raise (default) or null."#,
Some('n'),
)
.category(Category::Custom("dataframe".into()))
}
fn examples(&self) -> Vec<Example<'_>> {
vec![
Example {
description: "Apply timezone to a naive datetime",
example: r#"["2021-12-30 00:00:00" "2021-12-31 00:00:00"] | polars into-df
| polars as-datetime "%Y-%m-%d %H:%M:%S" --naive
| polars select (polars col datetime | polars replace-time-zone "America/New_York")"#,
result: Some(
NuDataFrame::try_from_columns(
vec![Column::new(
"datetime".to_string(),
vec![
Value::date(
DateTime::parse_from_str(
"2021-12-30 00:00:00 -0500",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
Value::date(
DateTime::parse_from_str(
"2021-12-31 00:00:00 -0500",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
],
)],
Some(NuSchema::new(Arc::new(Schema::from_iter(vec![
Field::new(
"datetime".into(),
DataType::Datetime(
TimeUnit::Nanoseconds,
TimeZone::opt_try_new(Some("America/New_York"))
.expect("timezone should be valid"),
),
),
])))),
Span::test_data(),
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
Example {
description: "Apply timezone with ambiguous datetime",
example: r#"["2025-11-02 00:00:00", "2025-11-02 01:00:00", "2025-11-02 02:00:00", "2025-11-02 03:00:00"]
| polars into-df
| polars as-datetime "%Y-%m-%d %H:%M:%S" --naive
| polars select (polars col datetime | polars replace-time-zone "America/New_York" --ambiguous null)"#,
result: Some(
NuDataFrame::try_from_columns(
vec![Column::new(
"datetime".to_string(),
vec![
Value::date(
DateTime::parse_from_str(
"2025-11-02 00:00:00 -0400",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
Value::nothing(Span::test_data()),
Value::date(
DateTime::parse_from_str(
"2025-11-02 02:00:00 -0500",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
Value::date(
DateTime::parse_from_str(
"2025-11-02 03:00:00 -0500",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
],
)],
Some(NuSchema::new(Arc::new(Schema::from_iter(vec![
Field::new(
"datetime".into(),
DataType::Datetime(
TimeUnit::Nanoseconds,
TimeZone::opt_try_new(Some("America/New_York"))
.expect("timezone should be valid"),
),
),
])))),
Span::test_data(),
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
Example {
description: "Apply timezone with nonexistent datetime",
example: r#"["2025-03-09 01:00:00", "2025-03-09 02:00:00", "2025-03-09 03:00:00", "2025-03-09 04:00:00"]
| polars into-df
| polars as-datetime "%Y-%m-%d %H:%M:%S" --naive
| polars select (polars col datetime | polars replace-time-zone "America/New_York" --nonexistent null)"#,
result: Some(
NuDataFrame::try_from_columns(
vec![Column::new(
"datetime".to_string(),
vec![
Value::date(
DateTime::parse_from_str(
"2025-03-09 01:00:00 -0500",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
Value::nothing(Span::test_data()),
Value::date(
DateTime::parse_from_str(
"2025-03-09 03:00:00 -0400",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
Value::date(
DateTime::parse_from_str(
"2025-03-09 04:00:00 -0400",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
],
)],
Some(NuSchema::new(Arc::new(Schema::from_iter(vec![
Field::new(
"datetime".into(),
DataType::Datetime(
TimeUnit::Nanoseconds,
TimeZone::opt_try_new(Some("America/New_York"))
.expect("timezone should be valid"),
),
),
])))),
Span::test_data(),
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
]
}
fn run(
&self,
plugin: &Self::Plugin,
engine: &EngineInterface,
call: &EvaluatedCall,
mut input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let metadata = input.take_metadata();
let value = input.into_value(call.head)?;
let ambiguous = match call.get_flag::<Value>("ambiguous")? {
Some(v @ Value::String { .. }) => {
let span = v.span();
let val = v.into_string()?;
match val.as_str() {
"raise" | "earliest" | "latest" => Ok(val),
_ => Err(ShellError::Generic(GenericError::new(
"Invalid argument value",
"`ambiguous` must be one of raise, earliest, latest, or null",
span,
))),
}
}
Some(Value::Nothing { .. }) => Ok("null".into()),
Some(_) => unreachable!("Argument only accepts string or null."),
None => Ok("raise".into()),
}
.map_err(LabeledError::from)?;
let nonexistent = match call.get_flag::<Value>("nonexistent")? {
Some(v @ Value::String { .. }) => match v.as_str()? {
"raise" => Ok(NonExistent::Raise),
_ => Err(ShellError::Generic(GenericError::new(
"Invalid argument value",
"`nonexistent` must be one of raise or null",
v.span(),
))),
},
Some(Value::Nothing { .. }) => Ok(NonExistent::Null),
Some(_) => unreachable!("Argument only accepts string or null."),
None => Ok(NonExistent::Raise),
}
.map_err(LabeledError::from)?;
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuExpression(expr) => {
command_expr(plugin, engine, call, expr, ambiguous, nonexistent)
}
PolarsPluginObject::NuSelector(selector) => {
let expr = selector.into_expr();
command_expr(plugin, engine, call, expr, ambiguous, nonexistent)
}
_ => Err(cant_convert_err(
&value,
&[PolarsPluginType::NuExpression, PolarsPluginType::NuSelector],
)),
}
.map_err(LabeledError::from)
.map(|pd| pd.set_metadata(metadata))
}
}
fn command_expr(
plugin: &PolarsPlugin,
engine: &EngineInterface,
call: &EvaluatedCall,
expr: NuExpression,
ambiguous: String,
nonexistent: NonExistent,
) -> Result<PipelineData, ShellError> {
let time_zone_spanned: Spanned<String> = call.req(0)?;
let time_zone = timezone_from_str(&time_zone_spanned.item, Some(time_zone_spanned.span))?;
let expr: NuExpression = expr
.into_polars()
.dt()
.replace_time_zone(
Some(time_zone),
Expr::Literal(LiteralValue::Dyn(DynLiteralValue::Str(
PlSmallStr::from_string(ambiguous),
))),
nonexistent,
)
.into();
expr.to_pipeline_data(plugin, engine, call.head)
}
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command;
use nu_protocol::ShellError;
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command(&ReplaceTimeZone)
}
}