From 09417bc6ec3aa74eef4b66764d9450459106399c Mon Sep 17 00:00:00 2001 From: MarcoGorelli <> Date: Sat, 28 Jan 2023 11:50:17 +0000 Subject: [PATCH 1/3] fix(rust, python): raise if tz_localize called on UTC-aware --- .../polars-plan/src/dsl/function_expr/datetime.rs | 5 +---- py-polars/tests/unit/test_datelike.py | 15 +++++++++++++++ 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs index 75b8e3fd627f..f7b657d4c343 100644 --- a/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs +++ b/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs @@ -146,10 +146,7 @@ pub(super) fn cast_timezone(s: &Series, tz: &str) -> PolarsResult { pub(super) fn tz_localize(s: &Series, tz: &str) -> PolarsResult { let ca = s.datetime()?.clone(); match ca.time_zone() { - Some(tz) if tz == "UTC" => { - Ok(ca.with_time_zone(Some("UTC".into()))?.into_series()) - } - Some(tz) if !tz.is_empty() => { + Some(old_tz) if !old_tz.is_empty() => { Err(PolarsError::ComputeError("Cannot localize a tz-aware datetime. Consider using 'dt.with_time_zone' or 'dt.cast_time_zone'".into())) }, _ => { diff --git a/py-polars/tests/unit/test_datelike.py b/py-polars/tests/unit/test_datelike.py index 62f0915bb198..b880b05dd04f 100644 --- a/py-polars/tests/unit/test_datelike.py +++ b/py-polars/tests/unit/test_datelike.py @@ -2181,6 +2181,21 @@ def test_tz_localize() -> None: } +@pytest.mark.parametrize("time_zone", ["UTC", "Africa/Abidjan"]) +def test_tz_localize_from_utc(time_zone: str) -> None: + ts_utc = ( + pl.Series(["2018-10-28"]).str.strptime(pl.Datetime).dt.tz_localize(time_zone) + ) + with pytest.raises( + ComputeError, + match=( + "^Cannot localize a tz-aware datetime. Consider using " + "'dt.with_time_zone' or 'dt.cast_time_zone'$" + ), + ): + ts_utc.dt.tz_localize("America/Maceio") + + def test_tz_aware_truncate() -> None: test = pl.DataFrame( { From 64484ec46506d4b5789ec8ba2868b8a15dc74620 Mon Sep 17 00:00:00 2001 From: MarcoGorelli <> Date: Sat, 28 Jan 2023 20:52:13 +0000 Subject: [PATCH 2/3] rewrite --- polars/polars-core/src/utils/supertype.rs | 6 +----- .../polars-plan/src/dsl/function_expr/datetime.rs | 9 ++++++--- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/polars/polars-core/src/utils/supertype.rs b/polars/polars-core/src/utils/supertype.rs index 909512cbc9e2..66885f8e4b41 100644 --- a/polars/polars-core/src/utils/supertype.rs +++ b/polars/polars-core/src/utils/supertype.rs @@ -206,11 +206,7 @@ pub fn get_supertype(l: &DataType, r: &DataType) -> Option { #[cfg(all(feature = "dtype-duration", feature = "dtype-datetime"))] (Duration(lu), Datetime(ru, Some(tz))) | (Datetime(lu, Some(tz)), Duration(ru)) => { - if tz.is_empty() { - Some(Datetime(get_time_units(lu, ru), None)) - } else { - Some(Datetime(get_time_units(lu, ru), Some(tz.clone()))) - } + Some(Datetime(get_time_units(lu, ru), Some(tz.clone()))) } #[cfg(all(feature = "dtype-duration", feature = "dtype-datetime"))] (Duration(lu), Datetime(ru, None)) | (Datetime(lu, None), Duration(ru)) => { diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs index f7b657d4c343..929bd160cd1e 100644 --- a/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs +++ b/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs @@ -145,11 +145,14 @@ pub(super) fn cast_timezone(s: &Series, tz: &str) -> PolarsResult { #[cfg(feature = "timezones")] pub(super) fn tz_localize(s: &Series, tz: &str) -> PolarsResult { let ca = s.datetime()?.clone(); - match ca.time_zone() { - Some(old_tz) if !old_tz.is_empty() => { + match (ca.time_zone(), tz) { + (Some(old_tz), _) if !old_tz.is_empty() => { Err(PolarsError::ComputeError("Cannot localize a tz-aware datetime. Consider using 'dt.with_time_zone' or 'dt.cast_time_zone'".into())) }, - _ => { + (_, "UTC") => { + Ok(ca.with_time_zone(Some("UTC".into()))?.into_series()) + } + (_, _) => { Ok(ca.with_time_zone(Some("UTC".to_string()))?.cast_time_zone(tz)?.into_series()) } } From c5a1f395cb9f57703082f8f393d513abde2ddb4c Mon Sep 17 00:00:00 2001 From: MarcoGorelli <> Date: Sat, 28 Jan 2023 20:54:15 +0000 Subject: [PATCH 3/3] revert unrelated --- polars/polars-core/src/utils/supertype.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/polars/polars-core/src/utils/supertype.rs b/polars/polars-core/src/utils/supertype.rs index 66885f8e4b41..909512cbc9e2 100644 --- a/polars/polars-core/src/utils/supertype.rs +++ b/polars/polars-core/src/utils/supertype.rs @@ -206,7 +206,11 @@ pub fn get_supertype(l: &DataType, r: &DataType) -> Option { #[cfg(all(feature = "dtype-duration", feature = "dtype-datetime"))] (Duration(lu), Datetime(ru, Some(tz))) | (Datetime(lu, Some(tz)), Duration(ru)) => { - Some(Datetime(get_time_units(lu, ru), Some(tz.clone()))) + if tz.is_empty() { + Some(Datetime(get_time_units(lu, ru), None)) + } else { + Some(Datetime(get_time_units(lu, ru), Some(tz.clone()))) + } } #[cfg(all(feature = "dtype-duration", feature = "dtype-datetime"))] (Duration(lu), Datetime(ru, None)) | (Datetime(lu, None), Duration(ru)) => {