Merge pull request #185 from HackSoc/improve-termdates

termdates: improve and refactor the plugin, fixing a couple of subtle bugs

Author: alanbriolat

src/csbot/plugins/termdates.py

@@ -1,43 +1,95 @@
 from csbot.plugin import Plugin
-from datetime import datetime, timedelta
+import datetime
 import math
+import typing as _t
 
 from ..util import ordinal
 
 
+class Term:
+    def __init__(self, key: str, start_date: datetime.datetime):
+        self.key = key
+        self.start_date = start_date
+
+    @property
+    def first_monday(self) -> datetime.datetime:
+        return self.start_date - datetime.timedelta(days=self.start_date.weekday())
+
+    @property
+    def last_friday(self) -> datetime.datetime:
+        return self.first_monday + datetime.timedelta(days=4, weeks=9)
+
+    def get_week_number(self, date: datetime.date) -> int:
+        """Get the "term week number" of a date relative to this term.
+
+        The first week of term is week 1, not week 0. Week 1 starts at the
+        Monday of the term's start date, even if the term's start date is not
+        Monday. Any date before the start of the term gives a negative week
+        number.
+        """
+        delta = date - self.first_monday.date()
+        week_number = math.floor(delta.days / 7.0)
+        if week_number >= 0:
+            return week_number + 1
+        else:
+            return week_number
+
+    def get_week_start(self, week_number: int) -> datetime.datetime:
+        """Get the start date of a specific week number relative to this term.
+
+        The first week of term is week 1, not week 0, although this method
+        allows both. When referring to the first week of term, the start date is
+        the term start date (which may not be a Monday). All other weeks start
+        on their Monday.
+        """
+        if week_number in (0, 1):
+            return self.start_date
+        elif week_number > 1:
+            return self.first_monday + datetime.timedelta(weeks=week_number - 1)
+        else:
+            return self.first_monday + datetime.timedelta(weeks=week_number)
+
+
 class TermDates(Plugin):
     """
     A wonderful plugin allowing old people (graduates) to keep track of the
     ever-changing calendar.
     """
     DATE_FORMAT = '%Y-%m-%d'
+    TERM_KEYS = ('aut', 'spr', 'sum')
 
     db_terms = Plugin.use('mongodb', collection='terms')
-    db_weeks = Plugin.use('mongodb', collection='weeks')
+
+    terms = None
+    _doc_id = None
 
     def setup(self):
         super(TermDates, self).setup()
+        self._load()
+
+    def _load(self):
+        doc = self.db_terms.find_one()
+        if not doc:
+            return False
+        self.terms = {key: Term(key, doc[key][0]) for key in self.TERM_KEYS}
+        self._doc_id = doc['_id']
+        return True
+
+    def _save(self):
+        if not self.terms:
+            return False
+        doc = {key: (self.terms[key].start_date, self.terms[key].last_friday) for key in self.TERM_KEYS}
+        if self._doc_id:
+            self.db_terms.replace_one({'_id': self._doc_id}, doc, upsert=True)
+        else:
+            res = self.db_terms.insert_one(doc)
+            self._doc_id = res.inserted_id
+        return True
 
-        # If we have stuff in mongodb, we can just load it directly.
-        if self.db_terms.find_one():
-            self.initialised = True
-            self.terms = self.db_terms.find_one()
-            self.weeks = self.db_weeks.find_one()
-            return
-
-        # If no term dates have been set, the calendar is uninitialised and
-        # can't be asked about term things.
-        self.initialised = False
-
-        # Each term is represented as a tuple of the date of the first Monday
-        # and the last Friday in it.
-        self.terms = {term: (None, None)
-                      for term in ['aut', 'spr', 'sum']}
-
-        # And each week is just the date of the Monday
-        self.weeks = {'{} {}'.format(term, week): None
-                      for term in ['aut', 'spr', 'sum']
-                      for week in range(1, 11)}
+    @property
+    def initialised(self) -> bool:
+        """If no term dates have been set, the calendar is uninitialised and can't be asked about term thing."""
+        return self._doc_id is not None
 
     @Plugin.command('termdates', help='termdates: show the current term dates')
     def termdates(self, e):
@@ -55,15 +107,15 @@ def _term_start(self, term):
         """
 
         term = term.lower()
-        return self.terms[term][0].strftime(self.DATE_FORMAT)
+        return self.terms[term].start_date.strftime(self.DATE_FORMAT)
 
     def _term_end(self, term):
         """
         Get the end date (last Friday) of a term as a string.
         """
 
         term = term.lower()
-        return self.terms[term][1].strftime(self.DATE_FORMAT)
+        return self.terms[term].last_friday.strftime(self.DATE_FORMAT)
 
     @Plugin.command('week',
                     help='week [term] [num]: info about a week, '
@@ -80,81 +132,74 @@ def week(self, e):
         #  !week term n - get the date of week n in the given term
         #  !week n term - as above
 
-        week = e['data'].split()
+        week = e['data'].lower().split()
         if len(week) == 0:
-            term, weeknum = self._current_week()
+            term, week_number = self._current_week()
         elif len(week) == 1:
             try:
-                term = self._current_term()
-                weeknum = int(week[0])
-                if weeknum < 1:
+                term = self._current_or_next_term()
+                week_number = int(week[0])
+                if week_number < 1:
                     e.reply('error: bad week format')
                     return
             except ValueError:
-                term = week[0][:3]
-                term, weeknum = self._current_week(term)
+                term_key = week[0][:3]
+                term, week_number = self._current_week(term_key)
         elif len(week) >= 2:
             try:
-                term = week[0][:3]
-                weeknum = int(week[1])
+                term_key = week[0][:3]
+                week_number = int(week[1])
             except ValueError:
                 try:
-                    term = week[1][:3]
-                    weeknum = int(week[0])
+                    term_key = week[1][:3]
+                    week_number = int(week[0])
                 except ValueError:
                     e.reply('error: bad week format')
                     return
+            try:
+                term = self.terms[term_key]
+            except KeyError:
+                e.reply('error: bad week format')
+                return
         else:
             e.reply('error: bad week format')
             return
 
-        if weeknum > 0:
-            e.reply('{} {}: {}'.format(term.capitalize(),
-                                       weeknum,
-                                       self._week_start(term, weeknum)))
+        if term is None:
+            e.reply('error: no term dates (see termdates.set)')
+        elif week_number > 0:
+            e.reply('{} {}: {}'.format(term.key.capitalize(),
+                                       week_number,
+                                       term.get_week_start(week_number).strftime(self.DATE_FORMAT)))
         else:
             e.reply('{} week before {} (starts {})'
-                    .format(ordinal(-weeknum),
-                            term.capitalize(),
-                            self._week_start(term, 1)))
+                    .format(ordinal(-week_number),
+                            term.key.capitalize(),
+                            term.start_date.strftime(self.DATE_FORMAT)))
 
-    def _current_term(self):
+    def _current_or_next_term(self) -> _t.Optional[Term]:
         """
         Get the name of the current term
         """
 
-        now = datetime.now().date()
-        for term in ['aut', 'spr', 'sum']:
-            dates = self.terms[term]
-            if now >= dates[0].date() and now <= dates[1].date():
+        now = datetime.datetime.now().date()
+        for key in self.TERM_KEYS:
+            term = self.terms[key]
+            if now < term.first_monday.date():
                 return term
-            elif now <= dates[0].date():
-                # We can do this because the terms are ordered
+            elif now <= term.last_friday.date():
                 return term
+        return None
 
-    def _current_week(self, term=None):
-        if term is None:
-            term = self._current_term()
-        start, _ = self.terms[term]
-        now = datetime.now()
-        delta = now.date() - start.date()
-        weeknum = math.floor(delta.days / 7.0)
-        if weeknum >= 0:
-            weeknum += 1
-        return term, weeknum
-
-    def _week_start(self, term, week):
-        """
-        Get the start date of a week as a string.
-        """
-
-        term = term.lower()
-        start = self.weeks['{} 1'.format(term)]
-        if week > 0:
-            offset = timedelta(weeks=week - 1)
+    def _current_week(self, key: _t.Optional[str] = None) -> (_t.Optional[Term], _t.Optional[int]):
+        if key:
+            term = self.terms.get(key.lower())
         else:
-            offset = timedelta(weeks=week)
-        return (start + offset).strftime(self.DATE_FORMAT)
+            term = self._current_or_next_term()
+        if term:
+            return term, term.get_week_number(datetime.date.today())
+        else:
+            return None, None
 
     @Plugin.command('termdates.set',
                     help='termdates.set <aut> <spr> <sum>: set the term dates')
@@ -165,47 +210,15 @@ def termdates_set(self, e):
             e.reply('error: all three dates must be provided')
             return
 
-        # Firstly compute the start and end dates of each term
-        for term, date in zip(['aut', 'spr', 'sum'], dates):
+        terms = {}
+        for key, date in zip(self.TERM_KEYS, dates):
             try:
-                term_start = datetime.strptime(date, self.DATE_FORMAT)
+                term_start = datetime.datetime.strptime(date, self.DATE_FORMAT)
             except ValueError:
                 e.reply('error: dates must be in %Y-%M-%d format.')
                 return
 
-            # Not all terms start on a monday, so we need to compute the "real"
-            # term start used in all the other calculations.
-            # Fortunately Monday is used as the start of the week in Python's
-            # datetime stuff, which makes this really simple.
-            real_start = term_start - timedelta(days=term_start.weekday())
-
-            # Log for informational purposes
-            if not term_start == real_start:
-                self.log.info('Computed real_start as {} (from {})'.format(
-                    repr(real_start), repr(term_start)))
-
-            term_end = real_start + timedelta(days=4, weeks=9)
-            self.terms[term] = (term_start, term_end)
-
-            # Then the start of each week
-            self.weeks['{} 1'.format(term)] = term_start
-            for week in range(2, 11):
-                week_start = real_start + timedelta(weeks=week-1)
-                self.weeks['{} {}'.format(term, week)] = week_start
-
-        # Save to the database. As we don't touch the _id attribute in this
-        # method, this will cause `save` to override the previously-loaded
-        # entry (if there is one).
-        if '_id' in self.terms:
-            self.db_terms.replace_one({'_id': self.terms['_id']}, self.terms, upsert=True)
-        else:
-            res = self.db_terms.insert_one(self.terms)
-            self.terms['_id'] = res.inserted_id
-        if '_id' in self.weeks:
-            self.db_weeks.replace_one({'_id': self.weeks['_id']}, self.weeks, upsert=True)
-        else:
-            res = self.db_weeks.insert_one(self.weeks)
-            self.weeks['_id'] = res.inserted_id
+            terms[key] = Term(key, term_start)
 
-        # Finally, we're initialised!
-        self.initialised = True
+        self.terms = terms
+        self._save()

tests/conftest.py

@@ -81,6 +81,30 @@ async def irc_client(request, event_loop, config_example_mode, irc_client_class,
     return client
 
 
+class LineMatcher:
+    def __init__(self, f, description):
+        self.f = f
+        self.description = description
+
+    def __call__(self, line):
+        return self.f(line)
+
+    def __repr__(self):
+        return self.description
+
+    @classmethod
+    def equals(cls, other):
+        return cls(lambda line: line == other, f"`line == {other!r}`")
+
+    @classmethod
+    def contains(cls, other):
+        return cls(lambda line: other in line, f"`{other!r} in line`")
+
+    @classmethod
+    def endswith(cls, other):
+        return cls(lambda line: line.endswith(other), f"`line.endswith({other!r})`")
+
+
 class IRCClientHelper:
     def __init__(self, irc_client):
         self.client = irc_client
@@ -155,21 +179,7 @@ def assert_sent(self, matchers, *, any_order=False, reset_mock=True):
         if reset_mock:
             self.client.send_line.reset_mock()
 
-
-class LineMatcher:
-    def __init__(self, f, description):
-        self.f = f
-        self.description = description
-
-    def __call__(self, line):
-        return self.f(line)
-
-    def __repr__(self):
-        return self.description
-
-    @classmethod
-    def equals(cls, other):
-        return cls(lambda line: line == other, f"`line == {other!r}`")
+    match_line = LineMatcher
 
 
 @pytest.fixture