diff options
| -rw-r--r-- | Lib/calendar.py | 19 | ||||
| -rw-r--r-- | Lib/test/test_calendar.py | 21 |
2 files changed, 30 insertions, 10 deletions
diff --git a/Lib/calendar.py b/Lib/calendar.py index 85baf2e..92149b7 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", @@ -176,22 +177,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 6dad058..6c7cdd1 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): |