fix(rust, python): write tz-aware datetimes to csv (pola-rs#6135)

Author: ritchie46

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::<chrono_tz::Tz>() {
-                                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::<chrono_tz::Tz>() {
+            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, "{{")?;

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;

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<u8>, 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);

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"