FIX: make error locator better

jklymak · jklymak · commit a6a96e33b7d4 · 2018-11-04T10:07:13.000-08:00
diff --git a/lib/matplotlib/dates.py b/lib/matplotlib/dates.py
@@ -1414,11 +1414,9 @@ def tick_values(self, vmin, vmax):
         vmin = vmin.replace(year=ymin, **self.replaced)
         try:
             ticks = [vmin.astimezone(self.tz)]
-        except ValueError:
+        except ValueError as e:
             raise ValueError('naive datetime objects cannot be used '
-                             'with matplotlib for python < 3.6; '
-                             'astimezone() cannot '
-                             'be applied to a naive datetime.')
+                             'with matplotlib for python < 3.6; ') from e
         while True:
             dt = ticks[-1]
             if dt.year >= ymax:
diff --git a/lib/matplotlib/tests/test_dates.py b/lib/matplotlib/tests/test_dates.py
@@ -507,11 +507,7 @@ def _create_auto_date_locator(date1, date2):
         assert list(map(str, mdates.num2date(locator()))) == expected
 
 
-@pytest.mark.pytz
-@pytest.mark.skipif(not __has_pytz(), reason="Requires pytz")
 def test_auto_date_locator_intmult_tz():
-    import pytz
-
     def _create_auto_date_locator(date1, date2, tz):
         locator = mdates.AutoDateLocator(interval_multiples=True, tz=tz)
         locator.create_dummy_axis()
@@ -572,8 +568,8 @@ def _create_auto_date_locator(date1, date2, tz):
                 ]
                )
 
-    tz = pytz.timezone('US/Pacific')
-    d1 = tz.localize(datetime.datetime(1997, 1, 1))
+    tz = dateutil.tz.gettz('Canada/Pacific')
+    d1 = datetime.datetime(1997, 1, 1, tzinfo=tz)
     for t_delta, expected in results:
         with rc_context({'_internal.classic_mode': False}):
             d2 = d1 + t_delta
@@ -726,6 +722,29 @@ def attach_tz(dt, zi):
     _test_rrulewrapper(attach_tz, pytz.timezone)
 
 
+@pytest.mark.pytz
+@pytest.mark.skipif(not __has_pytz(), reason="Requires pytz")
+def test_yearlocator_pytz():
+    import pytz
+
+    tz = pytz.timezone('America/New_York')
+    x = [datetime.datetime(2010, 1, 1).astimezone(tz)
+            + datetime.timedelta(i) for i in range(2000)]
+    locator = mdates.AutoDateLocator(interval_multiples=False, tz=tz)
+    locator.create_dummy_axis()
+    locator.set_view_interval(mdates.date2num(x[0]),
+                              mdates.date2num(x[-1]))
+
+    np.testing.assert_allclose([734138.20833333, 734503.20833333,
+                                734869.20833333, 735234.20833333,
+                                735599.20833333], locator())
+    expected = ['2011-01-01 00:00:00-05:00', '2012-01-01 00:00:00-05:00',
+                '2013-01-01 00:00:00-05:00', '2014-01-01 00:00:00-05:00',
+                '2015-01-01 00:00:00-05:00']
+    st = list(map(str, mdates.num2date(locator(), tz=tz)))
+    assert st == expected
+
+
 def test_DayLocator():
     with pytest.raises(ValueError):
         mdates.DayLocator(interval=-1)
