From 957b75699fd1ed2eb18e8f7fa6fee1062fabad68 Mon Sep 17 00:00:00 2001 From: Alexander Belopolsky Date: Tue, 27 Sep 2016 20:26:39 -0400 Subject: [PATCH] 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. --- 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 index 5244b8d1ee7..b5472f33066 100644 --- a/Lib/calendar.py +++ b/Lib/calendar.py @@ -8,6 +8,7 @@ set the first day of the week (0=Monday, 6=Sunday).""" import sys import datetime import locale as _locale +from itertools import repeat __all__ = ["IllegalMonthError", "IllegalWeekdayError", "setfirstweekday", "firstweekday", "isleap", "leapdays", "weekday", "monthrange", @@ -174,22 +175,20 @@ class Calendar(object): 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 index 80ed6325881..5de3ce802df 100644 --- a/Lib/test/test_calendar.py +++ b/Lib/test/test_calendar.py @@ -502,6 +502,27 @@ class CalendarTestCase(unittest.TestCase): # 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):