Selaa lähdekoodia

make astimezone call in localize more safe (#4324)

make astimezone call in localize more safe; with tests
Matt Davis 7 vuotta sitten
vanhempi
commit
976515108a
2 muutettua tiedostoa jossa 34 lisäystä ja 2 poistoa
  1. 5 1
      celery/utils/time.py
  2. 29 1
      t/unit/utils/test_time.py

+ 5 - 1
celery/utils/time.py

@@ -14,6 +14,7 @@ from kombu.utils.functional import reprcall
 from kombu.utils.objects import cached_property
 from pytz import AmbiguousTimeError, FixedOffset
 from pytz import timezone as _timezone
+from pytz import utc
 
 from celery.five import python_2_unicode_compatible, string_t
 
@@ -302,7 +303,10 @@ def make_aware(dt, tz):
 
 def localize(dt, tz):
     """Convert aware :class:`~datetime.datetime` to another timezone."""
-    dt = dt.astimezone(tz)
+    if is_naive(dt):  # Ensure timezone aware datetime
+        dt = make_aware(dt, tz)
+    if dt.tzinfo == utc:
+        dt = dt.astimezone(tz)  # Always safe to call astimezone on utc zones
     try:
         _normalize = tz.normalize
     except AttributeError:  # non-pytz tz

+ 29 - 1
t/unit/utils/test_time.py

@@ -177,7 +177,12 @@ class test_make_aware:
 class test_localize:
 
     def test_tz_without_normalize(self):
-        tz = tzinfo()
+        class tzz(tzinfo):
+
+            def utcoffset(self, dt):
+                return None  # Mock no utcoffset specified
+
+        tz = tzz()
         assert not hasattr(tz, 'normalize')
         assert localize(make_aware(datetime.utcnow(), tz), tz)
 
@@ -186,6 +191,9 @@ class test_localize:
         class tzz(tzinfo):
             raises = None
 
+            def utcoffset(self, dt):
+                return None
+
             def normalize(self, dt, **kwargs):
                 self.normalized = True
                 if self.raises and kwargs and kwargs.get('is_dst') is None:
@@ -209,6 +217,26 @@ class test_localize:
         assert tz3.normalized
         assert tz3.raised
 
+    def test_localize_changes_utc_dt(self):
+        now_utc_time = datetime.now(tz=pytz.utc)
+        local_tz = pytz.timezone('US/Eastern')
+        localized_time = localize(now_utc_time, local_tz)
+        assert localized_time == now_utc_time
+
+    def test_localize_aware_dt_idempotent(self):
+        t = (2017, 4, 23, 21, 36, 59, 0)
+        local_zone = pytz.timezone('America/New_York')
+        local_time = datetime(*t)
+        local_time_aware = datetime(*t, tzinfo=local_zone)
+        alternate_zone = pytz.timezone('America/Detroit')
+        localized_time = localize(local_time_aware, alternate_zone)
+        assert localized_time == local_time_aware
+        assert local_zone.utcoffset(
+            local_time) == alternate_zone.utcoffset(local_time)
+        localized_utc_offset = localized_time.tzinfo.utcoffset(local_time)
+        assert localized_utc_offset == alternate_zone.utcoffset(local_time)
+        assert localized_utc_offset == local_zone.utcoffset(local_time)
+
 
 @pytest.mark.parametrize('s,expected', [
     (999, 999),