[Python-checkins] cpython (merge 3.5 -> 3.6): Issue #28253: Fixed calendar functions for extreme months: 0001-01 and 9999-12.

[alexander.belopolsky]

https://hg.python.org/cpython/rev/cd384c4b441a
changeset:   104116:cd384c4b441a
branch:      3.6
parent:      104113:81f27d3ab214
parent:      104115:c439bce36bf2
user:        Alexander Belopolsky <alexander.belopolsky at gmail.com>
date:        Tue Sep 27 20:27:55 2016 -0400
summary:
  Issue #28253: Fixed calendar functions for extreme months: 0001-01 and 9999-12.

Methods itermonthdays() and itermonthdays2() are reimplemented so that they
don't call itermonthdates() which can cause datetime.date under/overflow.

files:
  Lib/calendar.py           |  19 +++++++++----------
  Lib/test/test_calendar.py |  21 +++++++++++++++++++++
  2 files changed, 30 insertions(+), 10 deletions(-)


diff --git a/Lib/calendar.py b/Lib/calendar.py
--- a/Lib/calendar.py
+++ b/Lib/calendar.py
@@ -8,6 +8,7 @@
 import sys
 import datetime
 import locale as _locale
+from itertools import repeat
 
 __all__ = ["IllegalMonthError", "IllegalWeekdayError", "setfirstweekday",
            "firstweekday", "isleap", "leapdays", "weekday", "monthrange",
@@ -176,22 +177,20 @@
         Like itermonthdates(), but will yield (day number, weekday number)
         tuples. For days outside the specified month the day number is 0.
         """
-        for date in self.itermonthdates(year, month):
-            if date.month != month:
-                yield (0, date.weekday())
-            else:
-                yield (date.day, date.weekday())
+        for i, d in enumerate(self.itermonthdays(year, month), self.firstweekday):
+            yield d, i % 7
 
     def itermonthdays(self, year, month):
         """
         Like itermonthdates(), but will yield day numbers. For days outside
         the specified month the day number is 0.
         """
-        for date in self.itermonthdates(year, month):
-            if date.month != month:
-                yield 0
-            else:
-                yield date.day
+        day1, ndays = monthrange(year, month)
+        days_before = (day1 - self.firstweekday) % 7
+        yield from repeat(0, days_before)
+        yield from range(1, ndays + 1)
+        days_after = (self.firstweekday - day1 - ndays) % 7
+        yield from repeat(0, days_after)
 
     def monthdatescalendar(self, year, month):
         """
diff --git a/Lib/test/test_calendar.py b/Lib/test/test_calendar.py
--- a/Lib/test/test_calendar.py
+++ b/Lib/test/test_calendar.py
@@ -502,6 +502,27 @@
         # see #15421
         list(calendar.Calendar().itermonthdates(datetime.MAXYEAR, 12))
 
+    def test_itermonthdays(self):
+        for firstweekday in range(7):
+            cal = calendar.Calendar(firstweekday)
+            # Test the extremes, see #28253 and #26650
+            for y, m in [(1, 1), (9999, 12)]:
+                days = list(cal.itermonthdays(y, m))
+                self.assertIn(len(days), (35, 42))
+        # Test a short month
+        cal = calendar.Calendar(firstweekday=3)
+        days = list(cal.itermonthdays(2001, 2))
+        self.assertEqual(days, list(range(1, 29)))
+
+    def test_itermonthdays2(self):
+        for firstweekday in range(7):
+            cal = calendar.Calendar(firstweekday)
+            # Test the extremes, see #28253 and #26650
+            for y, m in [(1, 1), (9999, 12)]:
+                days = list(cal.itermonthdays2(y, m))
+                self.assertEqual(days[0][1], firstweekday)
+                self.assertEqual(days[-1][1], (firstweekday - 1) % 7)
+
 
 class MonthCalendarTestCase(unittest.TestCase):
     def setUp(self):

-- 
Repository URL: https://hg.python.org/cpython