Skip to content

Commit

Permalink
ZoneInfo
Browse files Browse the repository at this point in the history
  • Loading branch information
@stinodego
stinodego committed Feb 28, 2024
1 parent e14824f commit e442b6b
Show file tree
Hide file tree
Showing 11 changed files with 129 additions and 154 deletions.
38 changes: 16 additions & 22 deletions py-polars/polars/utils/convert.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,28 +8,11 @@
from polars.dependencies import _ZONEINFO_AVAILABLE, zoneinfo

if TYPE_CHECKING:
import sys
from datetime import tzinfo
from decimal import Decimal

from polars.type_aliases import TimeUnit

# the below shenanigans with ZoneInfo are all to handle a
# typing issue in py < 3.9 while preserving lazy-loading
if sys.version_info >= (3, 9):
from zoneinfo import ZoneInfo
elif _ZONEINFO_AVAILABLE:
from backports.zoneinfo._zoneinfo import ZoneInfo

def get_zoneinfo(key: str) -> ZoneInfo: # noqa: D103
pass

else:

@lru_cache(None)
def get_zoneinfo(key: str) -> ZoneInfo: # noqa: D103
return zoneinfo.ZoneInfo(key)


SECONDS_PER_DAY = 86_400
SECONDS_PER_HOUR = 3_600
Expand Down Expand Up @@ -187,14 +170,25 @@ def to_py_datetime(

def _localize_datetime(dt: datetime, time_zone: str) -> datetime:
# zone info installation should already be checked
_tzinfo: ZoneInfo | tzinfo
try:
_tzinfo = get_zoneinfo(time_zone)
tz = string_to_zoneinfo(time_zone)
except zoneinfo.ZoneInfoNotFoundError:
# try fixed offset, which is not supported by ZoneInfo
_tzinfo = _parse_fixed_tz_offset(time_zone)
tz = _parse_fixed_tz_offset(time_zone)

return dt.astimezone(tz)


@lru_cache(None)
def string_to_zoneinfo(key: str) -> Any:
"""
Convert a time zone string to a Python ZoneInfo object.
return dt.astimezone(_tzinfo)
This is a simple wrapper for the zoneinfo.ZoneInfo constructor.
The wrapper is useful because zoneinfo is not available on Python 3.8
and the backports module may not be installed.
"""
return zoneinfo.ZoneInfo(key)


# cache here as we have a single tz per column
Expand All @@ -209,7 +203,7 @@ def _parse_fixed_tz_offset(offset: str) -> tzinfo:
# minutes, then we can construct:
# tzinfo=timezone(timedelta(hours=..., minutes=...))
except ValueError:
msg = f"offset: {offset!r} not understood"
msg = f"unexpected time zone offset: {offset!r}"
raise ValueError(msg) from None

return dt_offset.tzinfo # type: ignore[return-value]
Expand Down
14 changes: 5 additions & 9 deletions py-polars/tests/unit/constructors/test_constructors.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
from __future__ import annotations

import sys
from collections import OrderedDict, namedtuple
from datetime import date, datetime, time, timedelta, timezone
from decimal import Decimal
Expand All @@ -15,22 +14,19 @@

import polars as pl
from polars.datatypes import PolarsDataType, numpy_char_code_to_dtype
from polars.dependencies import _ZONEINFO_AVAILABLE, dataclasses, pydantic
from polars.dependencies import dataclasses, pydantic
from polars.exceptions import TimeZoneAwareConstructorWarning
from polars.testing import assert_frame_equal, assert_series_equal
from polars.utils._construction import type_hints

if TYPE_CHECKING:
from collections.abc import Callable

from polars.datatypes import PolarsDataType

if sys.version_info >= (3, 9):
from zoneinfo import ZoneInfo
elif _ZONEINFO_AVAILABLE:
# Import from submodule due to typing issue with backports.zoneinfo package:
# https://github.com/pganssle/zoneinfo/issues/125
from backports.zoneinfo._zoneinfo import ZoneInfo

from polars.datatypes import PolarsDataType
else:
from polars.utils.convert import string_to_zoneinfo as ZoneInfo


# -----------------------------------------------------------------------------------
Expand Down
9 changes: 3 additions & 6 deletions py-polars/tests/unit/dataframe/test_df.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,14 +27,11 @@
from polars.utils._construction import iterable_to_pydf

if TYPE_CHECKING:
from polars.type_aliases import JoinStrategy, UniqueKeepStrategy

if sys.version_info >= (3, 9):
from zoneinfo import ZoneInfo

from polars.type_aliases import JoinStrategy, UniqueKeepStrategy
else:
# Import from submodule due to typing issue with backports.zoneinfo package:
# https://github.com/pganssle/zoneinfo/issues/125
from backports.zoneinfo._zoneinfo import ZoneInfo
from polars.utils.convert import string_to_zoneinfo as ZoneInfo


def test_version() -> None:
Expand Down
150 changes: 75 additions & 75 deletions py-polars/tests/unit/datatypes/test_temporal.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@

from polars.type_aliases import Ambiguous, PolarsTemporalType, TimeUnit
else:
from polars.utils.convert import get_zoneinfo as ZoneInfo
from polars.utils.convert import string_to_zoneinfo as ZoneInfo


def test_fill_null() -> None:
Expand Down Expand Up @@ -1406,7 +1406,7 @@ def test_replace_time_zone() -> None:
@pytest.mark.parametrize(
("to_tz", "tzinfo"),
[
("America/Barbados", ZoneInfo(key="America/Barbados")),
("America/Barbados", ZoneInfo("America/Barbados")),
(None, None),
],
)
Expand All @@ -1430,7 +1430,7 @@ def test_strptime_with_tz() -> None:
.str.strptime(pl.Datetime("us", "Africa/Monrovia"))
.item()
)
assert result == datetime(2020, 1, 1, 3, tzinfo=ZoneInfo(key="Africa/Monrovia"))
assert result == datetime(2020, 1, 1, 3, tzinfo=ZoneInfo("Africa/Monrovia"))


@pytest.mark.parametrize(
Expand Down Expand Up @@ -1496,7 +1496,7 @@ def test_convert_time_zone_lazy_schema() -> None:
def test_convert_time_zone_on_tz_naive() -> None:
ts = pl.Series(["2020-01-01"]).str.strptime(pl.Datetime)
result = ts.dt.convert_time_zone("Asia/Kathmandu").item()
expected = datetime(2020, 1, 1, 5, 45, tzinfo=ZoneInfo(key="Asia/Kathmandu"))
expected = datetime(2020, 1, 1, 5, 45, tzinfo=ZoneInfo("Asia/Kathmandu"))
assert result == expected
result = (
ts.dt.replace_time_zone("UTC").dt.convert_time_zone("Asia/Kathmandu").item()
Expand Down Expand Up @@ -1582,8 +1582,8 @@ def test_replace_time_zone_from_naive() -> None:
pl.col("date").cast(pl.Datetime).dt.replace_time_zone("America/New_York")
).to_dict(as_series=False) == {
"date": [
datetime(2022, 1, 1, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 1, 2, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 1, 1, 0, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 1, 2, 0, 0, tzinfo=ZoneInfo("America/New_York")),
]
}

Expand Down Expand Up @@ -1854,22 +1854,22 @@ def test_tz_aware_truncate() -> None:
result = df.with_columns(pl.col("dt").dt.truncate("1d").alias("trunced"))
expected = {
"dt": [
datetime(2022, 11, 1, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 1, 12, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 2, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 2, 12, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 3, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 3, 12, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 4, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 1, 0, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 11, 1, 12, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 11, 2, 0, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 11, 2, 12, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 11, 3, 0, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 11, 3, 12, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 11, 4, 0, 0, tzinfo=ZoneInfo("America/New_York")),
],
"trunced": [
datetime(2022, 11, 1, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 1, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 2, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 2, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 3, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 3, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 4, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 1, 0, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 11, 1, 0, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 11, 2, 0, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 11, 2, 0, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 11, 3, 0, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 11, 3, 0, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 11, 4, 0, 0, tzinfo=ZoneInfo("America/New_York")),
],
}
assert result.to_dict(as_series=False) == expected
Expand Down Expand Up @@ -1900,34 +1900,34 @@ def test_tz_aware_truncate() -> None:
datetime(2022, 1, 1, 6, 0),
],
"UTC": [
datetime(2021, 12, 31, 23, 0, tzinfo=ZoneInfo(key="UTC")),
datetime(2022, 1, 1, 0, 0, tzinfo=ZoneInfo(key="UTC")),
datetime(2022, 1, 1, 1, 0, tzinfo=ZoneInfo(key="UTC")),
datetime(2022, 1, 1, 2, 0, tzinfo=ZoneInfo(key="UTC")),
datetime(2022, 1, 1, 3, 0, tzinfo=ZoneInfo(key="UTC")),
datetime(2022, 1, 1, 4, 0, tzinfo=ZoneInfo(key="UTC")),
datetime(2022, 1, 1, 5, 0, tzinfo=ZoneInfo(key="UTC")),
datetime(2022, 1, 1, 6, 0, tzinfo=ZoneInfo(key="UTC")),
datetime(2021, 12, 31, 23, 0, tzinfo=ZoneInfo("UTC")),
datetime(2022, 1, 1, 0, 0, tzinfo=ZoneInfo("UTC")),
datetime(2022, 1, 1, 1, 0, tzinfo=ZoneInfo("UTC")),
datetime(2022, 1, 1, 2, 0, tzinfo=ZoneInfo("UTC")),
datetime(2022, 1, 1, 3, 0, tzinfo=ZoneInfo("UTC")),
datetime(2022, 1, 1, 4, 0, tzinfo=ZoneInfo("UTC")),
datetime(2022, 1, 1, 5, 0, tzinfo=ZoneInfo("UTC")),
datetime(2022, 1, 1, 6, 0, tzinfo=ZoneInfo("UTC")),
],
"CST": [
datetime(2021, 12, 31, 17, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2021, 12, 31, 18, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2021, 12, 31, 19, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2021, 12, 31, 20, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2021, 12, 31, 21, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2021, 12, 31, 22, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2021, 12, 31, 23, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2022, 1, 1, 0, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2021, 12, 31, 17, 0, tzinfo=ZoneInfo("US/Central")),
datetime(2021, 12, 31, 18, 0, tzinfo=ZoneInfo("US/Central")),
datetime(2021, 12, 31, 19, 0, tzinfo=ZoneInfo("US/Central")),
datetime(2021, 12, 31, 20, 0, tzinfo=ZoneInfo("US/Central")),
datetime(2021, 12, 31, 21, 0, tzinfo=ZoneInfo("US/Central")),
datetime(2021, 12, 31, 22, 0, tzinfo=ZoneInfo("US/Central")),
datetime(2021, 12, 31, 23, 0, tzinfo=ZoneInfo("US/Central")),
datetime(2022, 1, 1, 0, 0, tzinfo=ZoneInfo("US/Central")),
],
"CST truncated": [
datetime(2021, 12, 31, 0, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2021, 12, 31, 0, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2021, 12, 31, 0, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2021, 12, 31, 0, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2021, 12, 31, 0, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2021, 12, 31, 0, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2021, 12, 31, 0, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2022, 1, 1, 0, 0, tzinfo=ZoneInfo(key="US/Central")),
datetime(2021, 12, 31, 0, 0, tzinfo=ZoneInfo("US/Central")),
datetime(2021, 12, 31, 0, 0, tzinfo=ZoneInfo("US/Central")),
datetime(2021, 12, 31, 0, 0, tzinfo=ZoneInfo("US/Central")),
datetime(2021, 12, 31, 0, 0, tzinfo=ZoneInfo("US/Central")),
datetime(2021, 12, 31, 0, 0, tzinfo=ZoneInfo("US/Central")),
datetime(2021, 12, 31, 0, 0, tzinfo=ZoneInfo("US/Central")),
datetime(2021, 12, 31, 0, 0, tzinfo=ZoneInfo("US/Central")),
datetime(2022, 1, 1, 0, 0, tzinfo=ZoneInfo("US/Central")),
],
}

Expand Down Expand Up @@ -1956,10 +1956,10 @@ def test_tz_aware_to_string() -> None:
result = df.with_columns(pl.col("dt").dt.to_string("%c").alias("fmt"))
expected = {
"dt": [
datetime(2022, 11, 1, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 2, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 3, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 4, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(2022, 11, 1, 0, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 11, 2, 0, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 11, 3, 0, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(2022, 11, 4, 0, 0, tzinfo=ZoneInfo("America/New_York")),
],
"fmt": [
"Tue Nov 1 00:00:00 2022",
Expand Down Expand Up @@ -2017,12 +2017,12 @@ def test_tz_aware_filter_lit() -> None:
datetime(1970, 1, 1, 5, 0),
],
"nyc": [
datetime(1970, 1, 1, 0, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(1970, 1, 1, 1, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(1970, 1, 1, 2, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(1970, 1, 1, 3, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(1970, 1, 1, 4, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(1970, 1, 1, 5, 0, tzinfo=ZoneInfo(key="America/New_York")),
datetime(1970, 1, 1, 0, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(1970, 1, 1, 1, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(1970, 1, 1, 2, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(1970, 1, 1, 3, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(1970, 1, 1, 4, 0, tzinfo=ZoneInfo("America/New_York")),
datetime(1970, 1, 1, 5, 0, tzinfo=ZoneInfo("America/New_York")),
],
}

Expand Down Expand Up @@ -2097,26 +2097,26 @@ def test_truncate_expr() -> None:
ambiguous_expr = df.select(pl.col("date").dt.truncate(every=pl.lit("30m")))
assert ambiguous_expr.to_dict(as_series=False) == {
"date": [
datetime(2020, 10, 25, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 0, 30, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 1, 0, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 1, 30, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 1, 0, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 1, 30, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 2, 0, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 0, 30, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 1, 0, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 1, 30, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 1, 0, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 1, 30, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 2, 0, tzinfo=ZoneInfo("Europe/London")),
]
}

all_expr = df.select(pl.col("date").dt.truncate(every=pl.col("every")))
assert all_expr.to_dict(as_series=False) == {
"date": [
datetime(2020, 10, 25, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 0, 45, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 1, 0, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 1, 45, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 1, 0, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 1, 45, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 2, 0, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 0, 45, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 1, 0, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 1, 45, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 1, 0, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 1, 45, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 2, 0, tzinfo=ZoneInfo("Europe/London")),
]
}

Expand Down Expand Up @@ -2352,13 +2352,13 @@ def test_round_ambiguous() -> None:
df = df.select(pl.col("date").dt.round("30m", ambiguous=pl.col("ambiguous")))
assert df.to_dict(as_series=False) == {
"date": [
datetime(2020, 10, 25, 0, 30, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 1, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 1, 30, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 1, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 1, 30, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 2, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 2, 30, tzinfo=ZoneInfo(key="Europe/London")),
datetime(2020, 10, 25, 0, 30, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 1, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 1, 30, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 1, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 1, 30, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 2, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 2, 30, tzinfo=ZoneInfo("Europe/London")),
]
}

Expand Down
15 changes: 6 additions & 9 deletions py-polars/tests/unit/expr/test_exprs.py
View file
Original file line number Diff line number Diff line change
@@ -1,16 +1,8 @@
from __future__ import annotations

import sys
from datetime import date, datetime, timedelta, timezone
from itertools import permutations
from typing import Any, cast

if sys.version_info >= (3, 9):
from zoneinfo import ZoneInfo
else:
# Import from submodule due to typing issue with backports.zoneinfo package:
# https://github.com/pganssle/zoneinfo/issues/125
from backports.zoneinfo._zoneinfo import ZoneInfo
from typing import TYPE_CHECKING, Any, cast

import numpy as np
import pytest
Expand All @@ -26,6 +18,11 @@
)
from polars.testing import assert_frame_equal, assert_series_equal

if TYPE_CHECKING:
from zoneinfo import ZoneInfo
else:
from polars.utils.convert import string_to_zoneinfo as ZoneInfo


def test_arg_true() -> None:
df = pl.DataFrame({"a": [1, 1, 2, 1]})
Expand Down
2 changes: 1 addition & 1 deletion py-polars/tests/unit/functions/as_datatype/test_as_datatype.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@

from polars.type_aliases import TimeUnit
else:
from polars.utils.convert import get_zoneinfo as ZoneInfo
from polars.utils.convert import string_to_zoneinfo as ZoneInfo


def test_date_datetime() -> None:
Expand Down
Loading

0 comments on commit e442b6b

Please sign in to comment.