diff --git a/sdks/python/apache_beam/utils/timestamp.py b/sdks/python/apache_beam/utils/timestamp.py index 81d8430d3af7..9ef900a26712 100644 --- a/sdks/python/apache_beam/utils/timestamp.py +++ b/sdks/python/apache_beam/utils/timestamp.py @@ -26,6 +26,7 @@ import datetime import re import time +import zoneinfo from typing import Optional from typing import Union from typing import overload @@ -166,7 +167,8 @@ def from_utc_datetime(cls, dt: datetime.datetime) -> 'Timestamp': "dt has no timezone info " + "(https://docs.python.org/3/library/datetime.html" + "#aware-and-naive-objects): %s" % dt) - if dt.tzinfo != pytz.utc and dt.tzinfo != datetime.timezone.utc: + if dt.tzinfo != pytz.utc and dt.tzinfo != datetime.timezone.utc and not ( + isinstance(dt.tzinfo, zoneinfo.ZoneInfo) and dt.tzinfo.key == 'UTC'): raise ValueError('dt not in UTC: %s' % dt) duration = dt - cls._epoch_datetime_utc() # Avoid total_seconds(): its float result can be off by a microsecond. diff --git a/sdks/python/apache_beam/utils/timestamp_test.py b/sdks/python/apache_beam/utils/timestamp_test.py index ebdd42d37029..fd004233398b 100644 --- a/sdks/python/apache_beam/utils/timestamp_test.py +++ b/sdks/python/apache_beam/utils/timestamp_test.py @@ -21,6 +21,7 @@ import datetime import unittest +import zoneinfo import pytz from google.protobuf import duration_pb2 @@ -86,6 +87,10 @@ def test_from_utc_datetime(self): Timestamp.from_utc_datetime( datetime.datetime(1970, 1, 1, tzinfo=pytz.utc)), Timestamp(0)) + self.assertEqual( + Timestamp.from_utc_datetime( + datetime.datetime(1970, 1, 1, tzinfo=zoneinfo.ZoneInfo("UTC"))), + Timestamp(0)) with self.assertRaisesRegex(ValueError, r'UTC'): Timestamp.from_utc_datetime( datetime.datetime(1970, 1, 1, tzinfo=pytz.timezone('US/Eastern')))