refactor: use standard library zoneinfo to replace pytz

Author: waketzheng

examples/fastapi/_tests.py

@@ -10,12 +10,12 @@
 
 import anyio
 import pytest
-import pytz
 from asgi_lifespan import LifespanManager
 from httpx import ASGITransport, AsyncClient
 
 from tortoise.contrib.test import MEMORY_SQLITE
 from tortoise.fields.data import JSON_LOADS
+from tortoise.timezone import UTC, parse_timezone
 
 os.environ["DB_URL"] = MEMORY_SQLITE
 try:
@@ -75,7 +75,7 @@ async def create_user(self, async_client: AsyncClient) -> Users:
         return user_obj
 
     async def user_list(self, async_client: AsyncClient) -> tuple[datetime, Users, User_Pydantic]:
-        utc_now = datetime.now(pytz.utc)
+        utc_now = datetime.now(UTC)
         user_obj = await Users.create(username="test")
         response = await async_client.get("/users")
         assert response.status_code == 200, response.text
@@ -123,13 +123,13 @@ async def test_create_user_east(self, client_east: AsyncClient) -> None:  # nose
         created_at = user_obj.created_at
 
         # Verify time zone
-        asia_tz = pytz.timezone(self.timezone)
-        asia_now = datetime.now(pytz.utc).astimezone(asia_tz)
+        asia_tz = parse_timezone(self.timezone)
+        asia_now = datetime.now(UTC).astimezone(asia_tz)
         assert created_at.hour - asia_now.hour == 0
 
         # UTC timezone
-        utc_tz = pytz.timezone("UTC")
-        utc_now = datetime.now(pytz.utc).astimezone(utc_tz)
+        utc_tz = parse_timezone("UTC")
+        utc_now = datetime.now(UTC).astimezone(utc_tz)
         assert (created_at.hour - utc_now.hour) in [self.delta_hours, self.delta_hours - 24]
 
     @pytest.mark.anyio

pyproject.toml

@@ -11,7 +11,6 @@ dependencies = [
     "pypika-tortoise (>=0.6.1,<1.0.0)",
     "iso8601 (>=2.1.0,<3.0.0); python_version < '4.0'",
     "aiosqlite (>=0.16.0,<1.0.0)",
-    "pytz",
 ]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",

tests/fields/test_time.py

@@ -4,7 +4,6 @@
 from time import sleep
 from unittest.mock import patch
 
-import pytz
 from iso8601 import ParseError
 
 from tests import testmodels
@@ -13,7 +12,7 @@
 from tortoise.contrib.test.condition import NotIn
 from tortoise.exceptions import ConfigurationError, IntegrityError
 from tortoise.expressions import F
-from tortoise.timezone import get_default_timezone
+from tortoise.timezone import get_default_timezone, parse_timezone
 
 
 class TestEmpty(test.TestCase):
@@ -126,7 +125,7 @@ async def test_set_timezone(self):
         old_tz = os.environ["TIMEZONE"]
         tz = "Asia/Shanghai"
         os.environ["TIMEZONE"] = tz
-        now = datetime.now(pytz.timezone(tz))
+        now = datetime.now(parse_timezone(tz))
         obj = await self.model.create(datetime=now)
         self.assertEqual(obj.datetime.tzinfo.zone, tz)
 
@@ -143,7 +142,7 @@ async def test_timezone(self):
         os.environ["TIMEZONE"] = tz
         os.environ["USE_TZ"] = "True"
 
-        now = datetime.now(pytz.timezone(tz))
+        now = datetime.now(parse_timezone(tz))
         obj = await self.model.create(datetime=now)
         self.assertEqual(obj.datetime.tzinfo.zone, tz)
         obj_get = await self.model.get(pk=obj.pk)

tests/test_default.py

@@ -1,8 +1,6 @@
 import datetime
 from decimal import Decimal
 
-import pytz
-
 from tests.testmodels import DefaultModel
 from tortoise import connections
 from tortoise.backends.asyncpg import AsyncpgDBClient
@@ -12,6 +10,7 @@
 from tortoise.backends.psycopg import PsycopgClient
 from tortoise.backends.sqlite import SqliteClient
 from tortoise.contrib import test
+from tortoise.timezone import UTC
 
 
 class TestDefault(test.TestCase):
@@ -45,5 +44,5 @@ async def test_default(self):
         self.assertEqual(default_model.date_default, datetime.date(year=2020, month=5, day=21))
         self.assertEqual(
             default_model.datetime_default,
-            datetime.datetime(year=2020, month=5, day=20, tzinfo=pytz.utc),
+            datetime.datetime(year=2020, month=5, day=20, tzinfo=UTC),
         )

tests/test_update.py

@@ -4,7 +4,6 @@
 from datetime import datetime, timedelta
 from typing import Any
 
-import pytz
 from pypika_tortoise.terms import Function as PupikaFunction
 
 from tests.testmodels import (
@@ -26,6 +25,7 @@
 from tortoise.contrib.test.condition import In, NotEQ
 from tortoise.expressions import Case, F, Q, Subquery, When
 from tortoise.functions import Function, Upper
+from tortoise.timezone import UTC
 
 
 class TestUpdate(test.TestCase):
@@ -49,11 +49,11 @@ async def test_bulk_update(self):
 
     async def test_bulk_update_datetime(self):
         objs = [
-            await DatetimeFields.create(datetime=datetime(2021, 1, 1, tzinfo=pytz.utc)),
-            await DatetimeFields.create(datetime=datetime(2021, 1, 1, tzinfo=pytz.utc)),
+            await DatetimeFields.create(datetime=datetime(2021, 1, 1, tzinfo=UTC)),
+            await DatetimeFields.create(datetime=datetime(2021, 1, 1, tzinfo=UTC)),
         ]
-        t0 = datetime(2021, 1, 2, tzinfo=pytz.utc)
-        t1 = datetime(2021, 1, 3, tzinfo=pytz.utc)
+        t0 = datetime(2021, 1, 2, tzinfo=UTC)
+        t1 = datetime(2021, 1, 3, tzinfo=UTC)
         objs[0].datetime = t0
         objs[1].datetime = t1
         rows_affected = await DatetimeFields.bulk_update(objs, fields=["datetime"])

tests/testmodels.py

@@ -13,7 +13,6 @@
 from enum import Enum, IntEnum
 from typing import Union
 
-import pytz
 from pydantic import BaseModel, ConfigDict
 
 from tortoise import fields
@@ -23,6 +22,7 @@
 from tortoise.manager import Manager
 from tortoise.models import Model
 from tortoise.queryset import QuerySet
+from tortoise.timezone import UTC
 from tortoise.validators import (
     CommaSeparatedIntegerListValidator,
     MaxValueValidator,
@@ -845,7 +845,7 @@ class DefaultModel(Model):
     char_default = fields.CharField(max_length=20, default="tortoise")
     date_default = fields.DateField(default=datetime.date(year=2020, month=5, day=21))
     datetime_default = fields.DatetimeField(
-        default=datetime.datetime(year=2020, month=5, day=20, tzinfo=pytz.utc)
+        default=datetime.datetime(year=2020, month=5, day=20, tzinfo=UTC)
     )
 
 

tortoise/backends/oracle/client.py

@@ -5,8 +5,6 @@
 from typing import TYPE_CHECKING, Any, SupportsInt, cast
 
 import pyodbc
-import pytz
-
 try:
     from ciso8601 import parse_datetime
 except ImportError:  # pragma: nocoverage
@@ -29,6 +27,7 @@
 )
 from tortoise.backends.oracle.executor import OracleExecutor
 from tortoise.backends.oracle.schema_generator import OracleSchemaGenerator
+from tortoise.timezone import UTC
 
 if TYPE_CHECKING:  # pragma: nocoverage
     import asyncodbc  # pylint: disable=W0611
@@ -102,7 +101,7 @@ def _timestamp_convert(self, value: bytes) -> datetime.date:
         try:
             return parse_datetime(value.decode()).date()
         except ValueError:
-            return parse_datetime(value.decode()[:-32]).astimezone(tz=pytz.utc)
+            return parse_datetime(value.decode()[:-32]).astimezone(tz=UTC)
 
     async def __aenter__(self) -> asyncodbc.Connection:
         connection = await super().__aenter__()

tortoise/timezone.py

@@ -2,9 +2,17 @@
 
 import functools
 import os
+import sys
 from datetime import datetime, time, tzinfo
+from zoneinfo import ZoneInfo as _ZoneInfo
+from zoneinfo import ZoneInfoNotFoundError
 
-import pytz
+if sys.version_info >= (3, 12):
+    from datetime import UTC
+else:
+    from datetime import timezone
+
+    UTC = timezone.utc
 
 
 @functools.cache
@@ -28,7 +36,7 @@ def now() -> datetime:
     Return an aware datetime.datetime, depending on use_tz and timezone.
     """
     if get_use_tz():
-        return datetime.now(tz=pytz.utc)
+        return datetime.now(tz=UTC)
     else:
         return datetime.now(get_default_timezone())
 
@@ -40,7 +48,32 @@ def get_default_timezone() -> tzinfo:
 
     This is the time zone defined by Tortoise config.
     """
-    return pytz.timezone(get_timezone())
+    return parse_timezone(get_timezone())
+
+
+class ZoneInfo(_ZoneInfo):
+    @property
+    def zone(self) -> str:
+        # Compatible with pytz:
+        # >>> ZoneInfo('UTC').key == pytz.timezone('UTC').zone == 'UTC'
+        return self.key
+
+
+def parse_timezone(zone: str) -> tzinfo:
+    if zone.upper() == "UTC":
+        return ZoneInfo("UTC")
+    try:
+        return ZoneInfo(zone)
+    except ZoneInfoNotFoundError as e:
+        words = zone.split("/")
+        # Compatible with `pytz.timezone`:
+        #   US/central -> US/Central
+        #   Europe/moscow -> Europe/Moscow
+        #   asia/ShangHai -> Asia/Shanghai
+        styled = "/".join([i if i.isupper() else i.title() for i in words])
+        if styled != zone:
+            return ZoneInfo(zone)
+        raise e
 
 
 def _reset_timezone_cache() -> None:
@@ -64,7 +97,7 @@ def localtime(value: datetime | None = None, timezone: str | None = None) -> dat
     """
     if value is None:
         value = now()
-    tz = get_default_timezone() if timezone is None else pytz.timezone(timezone)
+    tz = get_default_timezone() if timezone is None else parse_timezone(timezone)
     if is_naive(value):
         raise ValueError("localtime() cannot be applied to a naive datetime")
     return value.astimezone(tz)
@@ -104,7 +137,7 @@ def make_aware(
 
     :raises ValueError: when value is not naive datetime
     """
-    tz = get_default_timezone() if timezone is None else pytz.timezone(timezone)
+    tz = get_default_timezone() if timezone is None else parse_timezone(timezone)
     if hasattr(tz, "localize"):
         return tz.localize(value, is_dst=is_dst)
     if is_aware(value):
@@ -119,7 +152,7 @@ def make_naive(value: datetime, timezone: str | None = None) -> datetime:
 
     :raises ValueError: when value is naive datetime
     """
-    tz = get_default_timezone() if timezone is None else pytz.timezone(timezone)
+    tz = get_default_timezone() if timezone is None else parse_timezone(timezone)
     if is_naive(value):
         raise ValueError("make_naive() cannot be applied to a naive datetime")
     return value.astimezone(tz).replace(tzinfo=None)