Merge pull request #6262 from pganssle/fix_date2num_dst

tacaswell · web-flow · commit 0a31b3600787 · 2016-07-25T23:30:30.000-04:00
FIX: Properly handle UTC conversion in date2num
diff --git a/lib/matplotlib/dates.py b/lib/matplotlib/dates.py
@@ -212,47 +212,31 @@ def _to_ordinalf(dt):
     days, preserving hours, minutes, seconds and microseconds.  Return value
     is a :func:`float`.
     """
-
-    if hasattr(dt, 'tzinfo') and dt.tzinfo is not None:
-        delta = dt.tzinfo.utcoffset(dt)
-        if delta is not None:
-            dt -= delta
+    # Convert to UTC
+    tzi = getattr(dt, 'tzinfo', None)
+    if tzi is not None:
+        dt = dt.astimezone(UTC)
+        tzi = UTC
 
     base = float(dt.toordinal())
-    if isinstance(dt, datetime.datetime):
-        # Get a datetime object at midnight in the same time zone as dt.
-        cdate = dt.date()
-        midnight_time = datetime.time(0, 0, 0, tzinfo=dt.tzinfo)
+
+    # If it's sufficiently datetime-like, it will have a `date()` method
+    cdate = getattr(dt, 'date', lambda: None)()
+    if cdate is not None:
+        # Get a datetime object at midnight UTC
+        midnight_time = datetime.time(0, tzinfo=tzi)
 
         rdt = datetime.datetime.combine(cdate, midnight_time)
-        td_remainder = _total_seconds(dt - rdt)
 
-        if td_remainder > 0:
-            base += td_remainder / SEC_PER_DAY
+        # Append the seconds as a fraction of a day
+        base += (dt - rdt).total_seconds() / SEC_PER_DAY
 
     return base
 
 
 # a version of _to_ordinalf that can operate on numpy arrays
 _to_ordinalf_np_vectorized = np.vectorize(_to_ordinalf)
 
-try:
-    # Available as a native method in Python >= 2.7.
-    _total_seconds = datetime.timedelta.total_seconds
-except AttributeError:
-    def _total_seconds(tdelta):
-        """
-        Alias providing support for datetime.timedelta.total_seconds() function
-        calls even in Python < 2.7.
-
-        The input `tdelta` is a datetime.timedelta object, and returns a float
-        containing the total number of seconds representing the `tdelta`
-        duration. For large durations (> 270 on most platforms), this loses
-        microsecond accuracy.
-        """
-        return (tdelta.microseconds +
-                (tdelta.seconds + tdelta.days * SEC_PER_DAY) * 1e6) * 1e-6
-
 
 def _from_ordinalf(x, tz=None):
     """
@@ -432,7 +416,7 @@ def drange(dstart, dend, delta):
     """
     f1 = _to_ordinalf(dstart)
     f2 = _to_ordinalf(dend)
-    step = _total_seconds(delta) / SEC_PER_DAY
+    step = delta.total_seconds() / SEC_PER_DAY
 
     # calculate the difference between dend and dstart in times of delta
     num = int(np.ceil((f2 - f1) / step))
@@ -1065,8 +1049,8 @@ def get_locator(self, dmin, dmax):
         numDays = tdelta.days   # Avoids estimates of days/month, days/year
         numHours = (numDays * HOURS_PER_DAY) + delta.hours
         numMinutes = (numHours * MIN_PER_HOUR) + delta.minutes
-        numSeconds = np.floor(_total_seconds(tdelta))
-        numMicroseconds = np.floor(_total_seconds(tdelta) * 1e6)
+        numSeconds = np.floor(tdelta.total_seconds())
+        numMicroseconds = np.floor(tdelta.total_seconds() * 1e6)
 
         nums = [numYears, numMonths, numDays, numHours, numMinutes,
                 numSeconds, numMicroseconds]
@@ -1406,7 +1390,7 @@ def _close_to_dt(d1, d2, epsilon=5):
     Assert that datetimes *d1* and *d2* are within *epsilon* microseconds.
     """
     delta = d2 - d1
-    mus = abs(_total_seconds(delta) * 1e6)
+    mus = abs(delta.total_seconds() * 1e6)
     assert mus < epsilon
 
 
diff --git a/lib/matplotlib/tests/test_dates.py b/lib/matplotlib/tests/test_dates.py
@@ -9,12 +9,15 @@
 import tempfile
 
 import dateutil
+import pytz
+
 try:
     # mock in python 3.3+
     from unittest import mock
 except ImportError:
     import mock
 from nose.tools import assert_raises, assert_equal
+from nose.plugins.skip import SkipTest
 
 from matplotlib.testing.decorators import image_comparison, cleanup
 import matplotlib.pyplot as plt
@@ -355,6 +358,105 @@ def test_date_inverted_limit():
     fig.subplots_adjust(left=0.25)
 
 
+def _test_date2num_dst(date_range, tz_convert):
+    # Timezones
+    BRUSSELS = pytz.timezone('Europe/Brussels')
+    UTC = pytz.UTC
+
+    # Create a list of timezone-aware datetime objects in UTC
+    # Interval is 0b0.0000011 days, to prevent float rounding issues
+    dtstart = datetime.datetime(2014, 3, 30, 0, 0, tzinfo=UTC)
+    interval = datetime.timedelta(minutes=33, seconds=45)
+    interval_days = 0.0234375   # 2025 / 86400 seconds
+    N = 8
+
+    dt_utc = date_range(start=dtstart, freq=interval, periods=N)
+    dt_bxl = tz_convert(dt_utc, BRUSSELS)
+
+    expected_ordinalf = [735322.0 + (i * interval_days) for i in range(N)]
+    actual_ordinalf = list(mdates.date2num(dt_bxl))
+
+    assert_equal(actual_ordinalf, expected_ordinalf)
+
+
+def test_date2num_dst():
+    # Test for github issue #3896, but in date2num around DST transitions
+    # with a timezone-aware pandas date_range object.
+
+    class dt_tzaware(datetime.datetime):
+        """
+        This bug specifically occurs because of the normalization behavior of
+        pandas Timestamp objects, so in order to replicate it, we need a
+        datetime-like object that applies timezone normalization after
+        subtraction.
+        """
+        def __sub__(self, other):
+            r = super(dt_tzaware, self).__sub__(other)
+            tzinfo = getattr(r, 'tzinfo', None)
+
+            if tzinfo is not None:
+                localizer = getattr(tzinfo, 'normalize', None)
+                if localizer is not None:
+                    r = tzinfo.normalize(r)
+
+            if isinstance(r, datetime.datetime):
+                r = self.mk_tzaware(r)
+
+            return r
+
+        def __add__(self, other):
+            return self.mk_tzaware(super(dt_tzaware, self).__add__(other))
+
+        def astimezone(self, tzinfo):
+            dt = super(dt_tzaware, self).astimezone(tzinfo)
+            return self.mk_tzaware(dt)
+
+        @classmethod
+        def mk_tzaware(cls, datetime_obj):
+            kwargs = {}
+            attrs = ('year',
+                     'month',
+                     'day',
+                     'hour',
+                     'minute',
+                     'second',
+                     'microsecond',
+                     'tzinfo')
+
+            for attr in attrs:
+                val = getattr(datetime_obj, attr, None)
+                if val is not None:
+                    kwargs[attr] = val
+
+            return cls(**kwargs)
+
+    # Define a date_range function similar to pandas.date_range
+    def date_range(start, freq, periods):
+        dtstart = dt_tzaware.mk_tzaware(start)
+
+        return [dtstart + (i * freq) for i in range(periods)]
+
+    # Define a tz_convert function that converts a list to a new time zone.
+    def tz_convert(dt_list, tzinfo):
+        return [d.astimezone(tzinfo) for d in dt_list]
+
+    _test_date2num_dst(date_range, tz_convert)
+
+
+def test_date2num_dst_pandas():
+    # Test for github issue #3896, but in date2num around DST transitions
+    # with a timezone-aware pandas date_range object.
+    try:
+        import pandas as pd
+    except ImportError:
+        raise SkipTest('pandas not installed')
+
+    def tz_convert(*args):
+        return pd.DatetimeIndex.tz_convert(*args).astype(datetime.datetime)
+
+    _test_date2num_dst(pd.date_range, tz_convert)
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule(argv=['-s', '--with-doctest'], exit=False)
