Add 'zoneinfo' backend to TimezoneType (kvesteri#510)

This adds support for using the standard library's 'zoneinfo' module
as a backend for the TimezoneType. This module is available on Python
3.9+. For older versions, the backports.zoneinfo
module (https://github.com/pganssle/zoneinfo) is automatically used,
and is thus required on these versions.

Author: huonw

setup.py

@@ -33,6 +33,7 @@ def get_version():
         'pg8000>=1.12.4',
         'pytz>=2014.2',
         'python-dateutil>=2.6',
+        'backports.zoneinfo;python_version<"3.9"',
         'pymysql',
         'flake8>=2.4.0',
         'isort>=4.2.2',

sqlalchemy_utils/types/timezone.py

@@ -7,11 +7,9 @@
 
 class TimezoneType(ScalarCoercible, types.TypeDecorator):
     """
-    TimezoneType provides a way for saving timezones (from either the pytz or
-    the dateutil package) objects into database. TimezoneType saves timezone
-    objects as strings on the way in and converts them back to objects when
-    querying the database.
-
+    TimezoneType provides a way for saving timezones objects into database.
+    TimezoneType saves timezone objects as strings on the way in and converts
+    them back to objects when querying the database.
 
     ::
 
@@ -20,20 +18,22 @@ class TimezoneType(ScalarCoercible, types.TypeDecorator):
         class User(Base):
             __tablename__ = 'user'
 
-            # Pass backend='pytz' to change it to use pytz (dateutil by
-            # default)
+            # Pass backend='pytz' to change it to use pytz. Other values:
+            # 'dateutil' (default), and 'zoneinfo'.
             timezone = sa.Column(TimezoneType(backend='pytz'))
+
+    :param backend: Whether to use 'dateutil', 'pytz' or 'zoneinfo' for
+        timezones. 'zoneinfo' uses the standard library module in Python 3.9+,
+        but requires the external 'backports.zoneinfo' package for older
+        Python versions.
+
     """
 
     impl = types.Unicode(50)
 
     python_type = None
 
     def __init__(self, backend='dateutil'):
-        """
-        :param backend: Whether to use 'dateutil' or 'pytz' for timezones.
-        """
-
         self.backend = backend
         if backend == 'dateutil':
             try:
@@ -65,10 +65,27 @@ def __init__(self, backend='dateutil'):
                     "for 'TimezoneType'"
                 )
 
+        elif backend == "zoneinfo":
+            try:
+                import zoneinfo
+            except ImportError:
+                try:
+                    from backports import zoneinfo
+                except ImportError:
+                    raise ImproperlyConfigured(
+                        "'backports.zoneinfo' is required to use "
+                        "the 'zoneinfo' backend for 'TimezoneType'"
+                        "on Python version < 3.9"
+                    )
+
+            self.python_type = zoneinfo.ZoneInfo
+            self._to = zoneinfo.ZoneInfo
+            self._from = six.text_type
+
         else:
             raise ImproperlyConfigured(
-                "'pytz' or 'dateutil' are the backends supported for "
-                "'TimezoneType'"
+                "'pytz', 'dateutil' or 'zoneinfo' are the backends "
+                "supported for 'TimezoneType'"
             )
 
     def _coerce(self, value):

tests/types/test_timezone.py

@@ -3,6 +3,11 @@
 import sqlalchemy as sa
 from dateutil.zoneinfo import getzoneinfofile_stream, tzfile, ZoneInfoFile
 
+try:
+    import zoneinfo
+except ImportError:
+    from backports import zoneinfo
+
 from sqlalchemy_utils.types import timezone, TimezoneType
 
 
@@ -17,6 +22,9 @@ class Visitor(Base):
         timezone_pytz = sa.Column(
             timezone.TimezoneType(backend='pytz')
         )
+        timezone_zoneinfo = sa.Column(
+            timezone.TimezoneType(backend='zoneinfo')
+        )
 
         def __repr__(self):
             return 'Visitor(%r)' % self.id
@@ -32,7 +40,8 @@ class TestTimezoneType:
     def test_parameter_processing(self, session, Visitor):
         visitor = Visitor(
             timezone_dateutil=u'America/Los_Angeles',
-            timezone_pytz=u'America/Los_Angeles'
+            timezone_pytz=u'America/Los_Angeles',
+            timezone_zoneinfo=u'America/Los_Angeles'
         )
 
         session.add(visitor)
@@ -44,17 +53,21 @@ def test_parameter_processing(self, session, Visitor):
         visitor_pytz = session.query(Visitor).filter_by(
             timezone_pytz=u'America/Los_Angeles'
         ).first()
+        visitor_zoneinfo = session.query(Visitor).filter_by(
+            timezone_zoneinfo=u'America/Los_Angeles'
+        ).first()
 
         assert visitor_dateutil is not None
         assert visitor_pytz is not None
+        assert visitor_zoneinfo is not None
 
     def test_compilation(self, Visitor, session):
         query = sa.select([Visitor.timezone_pytz])
         # the type should be cacheable and not throw exception
         session.execute(query)
 
 
-TIMEZONE_BACKENDS = ['dateutil', 'pytz']
+TIMEZONE_BACKENDS = ['dateutil', 'pytz', 'zoneinfo']
 
 
 def test_can_coerce_pytz_DstTzInfo():
@@ -83,12 +96,23 @@ def test_can_coerce_string_for_dateutil_zone(zone):
     assert isinstance(tzcol._coerce(zone), tzfile)
 
 
+@pytest.mark.parametrize('zone', zoneinfo.available_timezones())
+def test_can_coerce_string_for_zoneinfo_zone(zone):
+    tzcol = TimezoneType(backend='zoneinfo')
+    assert str(tzcol._coerce(zone)) == zone
+
+
 @pytest.mark.parametrize('backend', TIMEZONE_BACKENDS)
 def test_can_coerce_and_raise_UnknownTimeZoneError_or_ValueError(backend):
     tzcol = TimezoneType(backend=backend)
-    with pytest.raises((ValueError, pytz.exceptions.UnknownTimeZoneError)):
+    exceptions = (
+        ValueError,
+        pytz.exceptions.UnknownTimeZoneError,
+        zoneinfo.ZoneInfoNotFoundError
+    )
+    with pytest.raises(exceptions):
         tzcol._coerce('SolarSystem/Mars')
-    with pytest.raises((ValueError, pytz.exceptions.UnknownTimeZoneError)):
+    with pytest.raises(exceptions):
         tzcol._coerce('')