feat(rust, python): avoid panic error in strftime with invalid format (…

…#6810) Co-authored-by: MarcoGorelli <>
pola-rs · Feb 12, 2023 · 194d9b8 · 194d9b8
1 parent d307963
commit 194d9b8
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 11 deletions.
diff --git a/polars/polars-core/src/chunked_array/temporal/datetime.rs b/polars/polars-core/src/chunked_array/temporal/datetime.rs
@@ -194,7 +194,7 @@ impl DatetimeChunked {
     }
 
     /// Format Datetime with a `fmt` rule. See [chrono strftime/strptime](https://docs.rs/chrono/0.4.19/chrono/format/strftime/index.html).
-    pub fn strftime(&self, fmt: &str) -> Utf8Chunked {
+    pub fn strftime(&self, fmt: &str) -> PolarsResult<Utf8Chunked> {
         #[cfg(feature = "timezones")]
         use chrono::Utc;
         let conversion_f = match self.time_unit() {
@@ -207,14 +207,26 @@ impl DatetimeChunked {
             .unwrap()
             .and_hms_opt(0, 0, 0)
             .unwrap();
-        let fmted = match self.time_zone() {
+        let mut fmted = String::new();
+        match self.time_zone() {
             #[cfg(feature = "timezones")]
-            Some(_) => format!(
+            Some(_) => write!(
+                fmted,
                 "{}",
                 Utc.from_local_datetime(&dt).earliest().unwrap().format(fmt)
-            ),
-            _ => format!("{}", dt.format(fmt)),
+            )
+            .map_err(|_| {
+                PolarsError::ComputeError(
+                    format!("Cannot format DateTime with format '{fmt}'.").into(),
+                )
+            })?,
+            _ => write!(fmted, "{}", dt.format(fmt)).map_err(|_| {
+                PolarsError::ComputeError(
+                    format!("Cannot format NaiveDateTime with format '{fmt}'.").into(),
+                )
+            })?,
         };
+        let fmted = fmted; // discard mut
 
         let mut ca: Utf8Chunked = match self.time_zone() {
             #[cfg(feature = "timezones")]
@@ -232,7 +244,7 @@ impl DatetimeChunked {
             _ => self.apply_kernel_cast(&|arr| format_naive(arr, fmt, &fmted, conversion_f)),
         };
         ca.rename(self.name());
-        ca
+        Ok(ca)
     }
 
     /// Construct a new [`DatetimeChunked`] from an iterator over [`NaiveDateTime`].

diff --git a/polars/polars-core/src/series/implementations/datetime.rs b/polars/polars-core/src/series/implementations/datetime.rs
@@ -342,13 +342,13 @@ impl SeriesTrait for SeriesWrap<DatetimeChunked> {
     fn cast(&self, data_type: &DataType) -> PolarsResult<Series> {
         match (data_type, self.0.time_unit()) {
             (DataType::Utf8, TimeUnit::Milliseconds) => {
-                Ok(self.0.strftime("%F %T%.3f").into_series())
+                Ok(self.0.strftime("%F %T%.3f")?.into_series())
             }
             (DataType::Utf8, TimeUnit::Microseconds) => {
-                Ok(self.0.strftime("%F %T%.6f").into_series())
+                Ok(self.0.strftime("%F %T%.6f")?.into_series())
             }
             (DataType::Utf8, TimeUnit::Nanoseconds) => {
-                Ok(self.0.strftime("%F %T%.9f").into_series())
+                Ok(self.0.strftime("%F %T%.9f")?.into_series())
             }
             _ => self.0.cast(data_type),
         }

diff --git a/polars/polars-time/src/series/mod.rs b/polars/polars-time/src/series/mod.rs
@@ -324,7 +324,9 @@ pub trait TemporalMethods: AsSeries {
             #[cfg(feature = "dtype-date")]
             DataType::Date => s.date().map(|ca| ca.strftime(fmt).into_series()),
             #[cfg(feature = "dtype-datetime")]
-            DataType::Datetime(_, _) => s.datetime().map(|ca| ca.strftime(fmt).into_series()),
+            DataType::Datetime(_, _) => {
+                s.datetime().map(|ca| Ok(ca.strftime(fmt)?.into_series()))?
+            }
             #[cfg(feature = "dtype-time")]
             DataType::Time => s.time().map(|ca| ca.strftime(fmt).into_series()),
             _ => Err(PolarsError::InvalidOperation(

diff --git a/py-polars/Cargo.lock b/py-polars/Cargo.lock
diff --git a/py-polars/tests/unit/test_datelike.py b/py-polars/tests/unit/test_datelike.py
@@ -1996,6 +1996,16 @@ def test_tz_aware_truncate() -> None:
     }
 
 
+def test_strftime_invalid_format() -> None:
+    tz_naive = pl.Series(["2020-01-01"]).str.strptime(pl.Datetime)
+    with pytest.raises(
+        ComputeError, match="Cannot format NaiveDateTime with format '%z'"
+    ):
+        tz_naive.dt.strftime("%z")
+    with pytest.raises(ComputeError, match="Cannot format DateTime with format '%q'"):
+        tz_naive.dt.replace_time_zone("UTC").dt.strftime("%q")
+
+
 def test_tz_aware_strftime() -> None:
     df = pl.DataFrame(
         {