Merge pull request #19348 from theOehrly/rrule-year-locator

timhoffm · web-flow · commit 3590c2634f38 · 2021-05-27T00:51:12.000+02:00
Make YearLocator a subclass of RRuleLocator
diff --git a/doc/api/next_api_changes/removals/19348-OE.rst b/doc/api/next_api_changes/removals/19348-OE.rst
@@ -0,0 +1,6 @@
+Removed ``dates.YearLocator.replaced`` attribute
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+`.YearLocator` is now a subclass of `.RRuleLocator`, and the attribute
+``YearLocator.replaced`` has been removed. For tick locations that
+required modifying this, a custom rrule and `.RRuleLocator` can be used instead.
diff --git a/lib/matplotlib/dates.py b/lib/matplotlib/dates.py
@@ -1168,6 +1168,15 @@ def __call__(self):
         return self.tick_values(dmin, dmax)
 
     def tick_values(self, vmin, vmax):
+        start, stop = self._create_rrule(vmin, vmax)
+        dates = self.rule.between(start, stop, True)
+        if len(dates) == 0:
+            return date2num([vmin, vmax])
+        return self.raise_if_exceeds(date2num(dates))
+
+    def _create_rrule(self, vmin, vmax):
+        # set appropriate rrule dtstart and until and return
+        # start and end
         delta = relativedelta(vmax, vmin)
 
         # We need to cap at the endpoints of valid datetime
@@ -1187,10 +1196,7 @@ def tick_values(self, vmin, vmax):
 
         self.rule.set(dtstart=start, until=stop)
 
-        dates = self.rule.between(vmin, vmax, True)
-        if len(dates) == 0:
-            return date2num([vmin, vmax])
-        return self.raise_if_exceeds(date2num(dates))
+        return vmin, vmax
 
     def _get_unit(self):
         # docstring inherited
@@ -1454,7 +1460,7 @@ def get_locator(self, dmin, dmax):
         return locator
 
 
-class YearLocator(DateLocator):
+class YearLocator(RRuleLocator):
     """
     Make ticks on a given day of each year that is a multiple of base.
 
@@ -1471,52 +1477,28 @@ def __init__(self, base=1, month=1, day=1, tz=None):
         Mark years that are multiple of base on a given month and day
         (default jan 1).
         """
-        super().__init__(tz)
+        rule = rrulewrapper(YEARLY, interval=base, bymonth=month,
+                            bymonthday=day, **self.hms0d)
+        super().__init__(rule, tz)
         self.base = ticker._Edge_integer(base, 0)
-        self.replaced = {'month':  month,
-                         'day':    day,
-                         'hour':   0,
-                         'minute': 0,
-                         'second': 0,
-                         }
-        if not hasattr(tz, 'localize'):
-            # if tz is pytz, we need to do this w/ the localize fcn,
-            # otherwise datetime.replace works fine...
-            self.replaced['tzinfo'] = tz
 
-    def __call__(self):
-        # if no data have been set, this will tank with a ValueError
-        try:
-            dmin, dmax = self.viewlim_to_dt()
-        except ValueError:
-            return []
+    def _create_rrule(self, vmin, vmax):
+        # 'start' needs to be a multiple of the interval to create ticks on
+        # interval multiples when the tick frequency is YEARLY
+        ymin = max(self.base.le(vmin.year) * self.base.step, 1)
+        ymax = min(self.base.ge(vmax.year) * self.base.step, 9999)
 
-        return self.tick_values(dmin, dmax)
+        c = self.rule._construct
+        replace = {'year': ymin,
+                   'month': c.get('bymonth', 1),
+                   'day': c.get('bymonthday', 1),
+                   'hour': 0, 'minute': 0, 'second': 0}
 
-    def tick_values(self, vmin, vmax):
-        ymin = self.base.le(vmin.year) * self.base.step
-        ymax = self.base.ge(vmax.year) * self.base.step
-
-        vmin = vmin.replace(year=ymin, **self.replaced)
-        if hasattr(self.tz, 'localize'):
-            # look after pytz
-            if not vmin.tzinfo:
-                vmin = self.tz.localize(vmin, is_dst=True)
-
-        ticks = [vmin]
-
-        while True:
-            dt = ticks[-1]
-            if dt.year >= ymax:
-                return date2num(ticks)
-            year = dt.year + self.base.step
-            dt = dt.replace(year=year, **self.replaced)
-            if hasattr(self.tz, 'localize'):
-                # look after pytz
-                if not dt.tzinfo:
-                    dt = self.tz.localize(dt, is_dst=True)
-
-            ticks.append(dt)
+        start = vmin.replace(**replace)
+        stop = start.replace(year=ymax)
+        self.rule.set(dtstart=start, until=stop)
+
+        return start, stop
 
 
 class MonthLocator(RRuleLocator):
diff --git a/lib/matplotlib/tests/test_dates.py b/lib/matplotlib/tests/test_dates.py
@@ -197,6 +197,18 @@ def test_RRuleLocator_dayrange():
     # On success, no overflow error shall be thrown
 
 
+def test_RRuleLocator_close_minmax():
+    # if d1 and d2 are very close together, rrule cannot create
+    # reasonable tick intervals; ensure that this is handled properly
+    rrule = mdates.rrulewrapper(dateutil.rrule.SECONDLY, interval=5)
+    loc = mdates.RRuleLocator(rrule)
+    d1 = datetime.datetime(year=2020, month=1, day=1)
+    d2 = datetime.datetime(year=2020, month=1, day=1, microsecond=1)
+    expected = ['2020-01-01 00:00:00+00:00',
+                '2020-01-01 00:00:00.000001+00:00']
+    assert list(map(str, mdates.num2date(loc.tick_values(d1, d2)))) == expected
+
+
 @image_comparison(['DateFormatter_fractionalSeconds.png'])
 def test_DateFormatter():
     import matplotlib.testing.jpl_units as units
@@ -964,6 +976,45 @@ def test_yearlocator_pytz():
     assert st == expected
 
 
+def test_YearLocator():
+    def _create_year_locator(date1, date2, **kwargs):
+        locator = mdates.YearLocator(**kwargs)
+        locator.create_dummy_axis()
+        locator.axis.set_view_interval(mdates.date2num(date1),
+                                       mdates.date2num(date2))
+        return locator
+
+    d1 = datetime.datetime(1990, 1, 1)
+    results = ([datetime.timedelta(weeks=52 * 200),
+                {'base': 20, 'month': 1, 'day': 1},
+                ['1980-01-01 00:00:00+00:00', '2000-01-01 00:00:00+00:00',
+                 '2020-01-01 00:00:00+00:00', '2040-01-01 00:00:00+00:00',
+                 '2060-01-01 00:00:00+00:00', '2080-01-01 00:00:00+00:00',
+                 '2100-01-01 00:00:00+00:00', '2120-01-01 00:00:00+00:00',
+                 '2140-01-01 00:00:00+00:00', '2160-01-01 00:00:00+00:00',
+                 '2180-01-01 00:00:00+00:00', '2200-01-01 00:00:00+00:00']
+                ],
+               [datetime.timedelta(weeks=52 * 200),
+                {'base': 20, 'month': 5, 'day': 16},
+                ['1980-05-16 00:00:00+00:00', '2000-05-16 00:00:00+00:00',
+                 '2020-05-16 00:00:00+00:00', '2040-05-16 00:00:00+00:00',
+                 '2060-05-16 00:00:00+00:00', '2080-05-16 00:00:00+00:00',
+                 '2100-05-16 00:00:00+00:00', '2120-05-16 00:00:00+00:00',
+                 '2140-05-16 00:00:00+00:00', '2160-05-16 00:00:00+00:00',
+                 '2180-05-16 00:00:00+00:00', '2200-05-16 00:00:00+00:00']
+                ],
+               [datetime.timedelta(weeks=52 * 5),
+                {'base': 20, 'month': 9, 'day': 25},
+                ['1980-09-25 00:00:00+00:00', '2000-09-25 00:00:00+00:00']
+                ],
+               )
+
+    for delta, arguments, expected in results:
+        d2 = d1 + delta
+        locator = _create_year_locator(d1, d2, **arguments)
+        assert list(map(str, mdates.num2date(locator()))) == expected
+
+
 def test_DayLocator():
     with pytest.raises(ValueError):
         mdates.DayLocator(interval=-1)
