diff --git a/polars/polars-core/src/fmt.rs b/polars/polars-core/src/fmt.rs index 2f9d63d992ca..4b357f625df3 100644 --- a/polars/polars-core/src/fmt.rs +++ b/polars/polars-core/src/fmt.rs @@ -9,6 +9,7 @@ use std::fmt::{Debug, Display, Formatter}; feature = "dtype-time" ))] use arrow::temporal_conversions::*; +use chrono::NaiveDateTime; #[cfg(feature = "timezones")] use chrono::TimeZone; #[cfg(any(feature = "fmt", feature = "fmt_no_tty"))] @@ -714,28 +715,8 @@ impl Display for AnyValue<'_> { }; match tz { None => write!(f, "{ndt}"), - Some(_tz) => { - #[cfg(feature = "timezones")] - { - match _tz.parse::() { - Ok(tz) => { - let dt_utc = chrono::Utc.from_local_datetime(&ndt).unwrap(); - let dt_tz_aware = dt_utc.with_timezone(&tz); - write!(f, "{dt_tz_aware}") - } - Err(_) => match parse_offset(_tz) { - Ok(offset) => { - let dt_tz_aware = offset.from_utc_datetime(&ndt); - write!(f, "{dt_tz_aware}") - } - Err(_) => write!(f, "invalid timezone"), - }, - } - } - #[cfg(not(feature = "timezones"))] - { - panic!("activate 'timezones' feature") - } + Some(tz) => { + write!(f, "{}", PlTzAware::new(ndt, tz)) } } } @@ -776,6 +757,43 @@ impl Display for AnyValue<'_> { } } +/// Utility struct to format a timezone aware datetime. +#[allow(dead_code)] +pub struct PlTzAware<'a> { + ndt: NaiveDateTime, + tz: &'a str, +} +impl<'a> PlTzAware<'a> { + pub fn new(ndt: NaiveDateTime, tz: &'a str) -> Self { + Self { ndt, tz } + } +} + +impl Display for PlTzAware<'_> { + #[allow(unused_variables)] + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + #[cfg(feature = "timezones")] + match self.tz.parse::() { + Ok(tz) => { + let dt_utc = chrono::Utc.from_local_datetime(&self.ndt).unwrap(); + let dt_tz_aware = dt_utc.with_timezone(&tz); + write!(f, "{dt_tz_aware}") + } + Err(_) => match parse_offset(self.tz) { + Ok(offset) => { + let dt_tz_aware = offset.from_utc_datetime(&self.ndt); + write!(f, "{dt_tz_aware}") + } + Err(_) => write!(f, "invalid timezone"), + }, + } + #[cfg(not(feature = "timezones"))] + { + panic!("activate 'timezones' feature") + } + } +} + #[cfg(feature = "dtype-struct")] fn fmt_struct(f: &mut Formatter<'_>, vals: &[AnyValue]) -> fmt::Result { write!(f, "{{")?; diff --git a/polars/polars-core/src/lib.rs b/polars/polars-core/src/lib.rs index 59edaf7d8527..238df3fb4b07 100644 --- a/polars/polars-core/src/lib.rs +++ b/polars/polars-core/src/lib.rs @@ -10,7 +10,7 @@ pub mod datatypes; pub mod doc; pub mod error; pub mod export; -mod fmt; +pub mod fmt; pub mod frame; pub mod functions; mod named_from; diff --git a/polars/polars-io/src/csv/write_impl.rs b/polars/polars-io/src/csv/write_impl.rs index 8adf34684f0c..7fc3e5123b84 100644 --- a/polars/polars-io/src/csv/write_impl.rs +++ b/polars/polars-io/src/csv/write_impl.rs @@ -4,6 +4,7 @@ use arrow::temporal_conversions; use lexical_core::{FormattedSize, ToLexical}; use memchr::{memchr, memchr2}; use polars_core::error::PolarsError::ComputeError; +use polars_core::fmt::PlTzAware; use polars_core::prelude::*; use polars_core::series::SeriesIter; use polars_core::POOL; @@ -83,28 +84,22 @@ fn write_anyvalue(f: &mut Vec, value: AnyValue, options: &SerializeOptions) } } #[cfg(feature = "dtype-datetime")] - AnyValue::Datetime(v, tu, tz) => match tz { - None => { - let dt = match tu { - TimeUnit::Nanoseconds => temporal_conversions::timestamp_ns_to_datetime(v), - TimeUnit::Microseconds => temporal_conversions::timestamp_us_to_datetime(v), - TimeUnit::Milliseconds => temporal_conversions::timestamp_ms_to_datetime(v), - }; - match &options.datetime_format { - None => write!(f, "{dt}"), - Some(fmt) => write!(f, "{}", dt.format(fmt)), - } - } - Some(tz) => { - let tz = temporal_conversions::parse_offset(tz).unwrap(); - - let dt = temporal_conversions::timestamp_to_datetime(v, tu.to_arrow(), &tz); - match &options.datetime_format { - None => write!(f, "{dt}"), - Some(fmt) => write!(f, "{}", dt.format(fmt)), + AnyValue::Datetime(v, tu, tz) => { + let ndt = match tu { + TimeUnit::Nanoseconds => temporal_conversions::timestamp_ns_to_datetime(v), + TimeUnit::Microseconds => temporal_conversions::timestamp_us_to_datetime(v), + TimeUnit::Milliseconds => temporal_conversions::timestamp_ms_to_datetime(v), + }; + match tz { + None => match &options.datetime_format { + None => write!(f, "{ndt}"), + Some(fmt) => write!(f, "{}", ndt.format(fmt)), + }, + Some(tz) => { + write!(f, "{}", PlTzAware::new(ndt, tz)) } } - }, + } #[cfg(feature = "dtype-time")] AnyValue::Time(v) => { let date = temporal_conversions::time64ns_to_time(v); diff --git a/py-polars/tests/unit/io/test_csv.py b/py-polars/tests/unit/io/test_csv.py index 333b0f902f41..b59d5514c91e 100644 --- a/py-polars/tests/unit/io/test_csv.py +++ b/py-polars/tests/unit/io/test_csv.py @@ -984,3 +984,10 @@ def test_csv_quoted_missing() -> None: "col3": [123, None, 101112], "col4": [456, 789, 131415], } + + +def test_csv_write_tz_aware() -> None: + df = pl.DataFrame({"times": datetime(2021, 1, 1)}).with_columns( + pl.col("times").dt.with_time_zone("Europe/Zurich") + ) + assert df.write_csv() == "times\n2021-01-01 01:00:00 CET\n"