Add docstrings and mypy annotations

Author: jmymay

.github/workflows/main.yml

@@ -27,6 +27,9 @@ jobs:
     - name: Lint with black
       run: |
         poetry run black --check .
+    - name: Check with mypy
+      run: |
+        poetry run mypy .
     - name: Run tests
       run: |
         poetry run pytest tests

.gitignore

@@ -4,3 +4,4 @@
 config.json
 *.egg-info
 __pycache__
+.mypy_cache

.pre-commit-config.yaml

@@ -16,3 +16,7 @@ repos:
     rev: 3.7.9
     hooks:
     -   id: flake8
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v0.761
+    hooks:
+    -   id: mypy

mypy.ini

@@ -0,0 +1,11 @@
+[mypy]
+python_version = 3.8
+
+[mypy-discord]
+ignore_missing_imports = true
+
+[mypy-pylru]
+ignore_missing_imports = true
+
+[mypy-pytest]
+ignore_missing_imports = true

oumodulesbot/backend.py

@@ -2,11 +2,14 @@
 import json
 import logging
 import re
+from typing import Mapping, Optional, Tuple
 
 from .ou_utils import MODULE_CODE_RE_TEMPLATE, get_module_url
 
 logger = logging.getLogger(__name__)
 
+CacheItem = Tuple[str, Optional[str]]  # title, url
+
 
 class OUModulesBackend:
 
@@ -17,9 +20,18 @@ class OUModulesBackend:
 
     def __init__(self):
         with open("cache.json", "r") as f:
-            self.cache = json.load(f)
+            cache_json = json.load(f)
+        self.cache: Mapping[str, CacheItem] = {
+            k: tuple(v) for k, v in cache_json.items()
+        }
+
+    async def get_module_title(self, code: str) -> Optional[str]:
+        """
+        Returns a module title for given code, if available.
 
-    async def get_module_title(self, code):
+        Tries a lookup in cache, and if it fails then it attempts to query
+        the Open University Digital Archive.
+        """
         code = code.upper()
 
         # 1. Try cached title:
@@ -36,17 +48,37 @@ async def get_module_title(self, code):
                 response = await client.get(url_template.format(code))
             html = response.content.decode("utf-8")
         except Exception:
-            return
+            return None
         title = self.MODULE_TITLE_RE.findall(html)
         return title[0].replace("!", "") if title else None
 
-    async def _is_module_url(self, url, code):
+    async def _is_module_url(self, url: str, code: str) -> bool:
+        """
+        Check if given URL looks like a valid URL for a given module code,
+        i.e. resolves to 200, and doesn't redirect away to a different page.
+
+        OU redirects to places like /courses/ sometimes which is a masked way
+        of saying '404'. However 301 redirects don't always indicate that
+        modules aren't available, because they sometimes point to a different
+        page for the same module.
+
+        Thus a compromise is used here by allowing redirects, but only if the
+        destination page URL includes the module code.
+        """
         async with httpx.AsyncClient() as client:
             response = await client.head(url, allow_redirects=True)
             correct_redirect = code.lower() in str(response.url).lower()
             return correct_redirect and response.status_code == 200
 
-    async def get_module_url(self, code):
+    async def get_module_url(self, code: str) -> Optional[str]:
+        """
+        Return module's URL for given module code, if available.
+
+        Tries to lookup in cache, and if it fails then it tries to provide
+        the URL with a well-known template (see get_module_url from ou_utils).
+
+        The template-based URL is returned only if it passes a HTTP check.
+        """
         code = code.upper()
 
         # 1. Try cached URL

oumodulesbot/ou_utils.py

@@ -1,13 +1,14 @@
 MODULE_CODE_RE_TEMPLATE = r"[a-zA-Z]{1,6}[0-9]{1,3}(?:-[a-zA-Z]{1,5})?"
 
 
-def get_module_level(module_code):
+def get_module_level(module_code: str) -> int:
     for c in module_code:
         if c.isdigit():
             return int(c)
+    raise ValueError(f"Invalid module code: {module_code}")
 
 
-def get_module_url(module_code):
+def get_module_url(module_code: str) -> str:
     if get_module_level(module_code) == 0:
         template = "http://www.open.ac.uk/courses/short-courses/{}"
     elif get_module_level(module_code) == 8:

oumodulesbot/oumodulesbot.py

@@ -1,3 +1,5 @@
+from collections import namedtuple
+from typing import Iterable, List, Sequence
 import json
 import logging
 import os
@@ -12,6 +14,7 @@
 logger = logging.getLogger(__name__)
 
 replies_cache = pylru.lrucache(1000)
+Module = namedtuple("Module", "code,title")
 
 
 class OUModulesBot(discord.Client):
@@ -23,8 +26,12 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.backend = OUModulesBackend()
 
-    async def do_mentions(self, message):
-        modules = []
+    async def process_mentions(self, message: discord.Message) -> None:
+        """
+        Process module code mentions from given `message`, and reply with
+        thieir names/URLs if any were found.
+        """
+        modules: List[Module] = []
         any_found = False
         for module in self.MENTION_RE.findall(message.content)[
             : self.MODULES_COUNT_LIMIT
@@ -33,14 +40,23 @@ async def do_mentions(self, message):
             title = await self.backend.get_module_title(module_code)
             if title:
                 any_found = True
-                modules.append((module_code, title))
+                modules.append(Module(module_code, title))
             else:
-                modules.append((module_code, "not found"))
+                modules.append(Module(module_code, "not found"))
         if any_found:
             # don't spam unless we're sure we at least found some modules
             await self.post_modules(message, modules)
 
-    async def format_course(self, code, title, for_embed=False):
+    async def format_course(
+        self, code: str, title: str, for_embed: bool = False
+    ) -> str:
+        """
+        Return a string describing a module ready for posting to Discord,
+        for given module `code` and `title`. Appends URL if available.
+
+        Uses more compact formatting if `for_embed` is True, which should
+        be used if multiple modules are presented as part of an embed.
+        """
         fmt = " * {} " if for_embed else "{}"
         fmt_link = " * [{}]({}) " if for_embed else "{} ({})"
         url = await self.backend.get_module_url(code)
@@ -53,24 +69,38 @@ async def format_course(self, code, title, for_embed=False):
         else:
             return "{}: {}".format(code, result)
 
-    async def _embed_modules(self, embed, modules):
+    async def embed_modules(
+        self, embed: discord.Embed, modules: Iterable[Module]
+    ) -> None:
+        """
+        Adds `embed` fields for each provided module.
+        """
         for (code, title) in modules:
             embed.add_field(
                 name=code,
                 value=await self.format_course(code, title, for_embed=True),
                 inline=True,
             )
 
-    async def post_modules(self, message, modules):
+    async def post_modules(
+        self, message: discord.Message, modules: Sequence[Module]
+    ) -> None:
+        """
+        Create or update a bot message for given users's input message,
+        and a list of modules.
+
+        Message is updated instead of created if the input was already replied
+        to, which means this time the input was edited.
+        """
         modify_message = None
         if message.id in replies_cache:
             modify_message = replies_cache[message.id]
 
         embed = discord.Embed()
         if len(modules) > 1:
             content = " "  # force removal when modifying
-            await self._embed_modules(embed, modules)
-        elif len(modules) > 0:
+            await self.embed_modules(embed, modules)
+        elif len(modules) == 1:
             code, title = modules[0]
             content = await self.format_course(code, title)
         else:
@@ -88,11 +118,13 @@ async def post_modules(self, message, modules):
                 content, embed=embed if len(modules) > 1 else None
             )
 
-    async def on_message(self, message):
-        await self.do_mentions(message)
+    async def on_message(self, message: discord.Message) -> None:
+        await self.process_mentions(message)
 
-    async def on_message_edit(self, before, after):
-        await self.do_mentions(after)
+    async def on_message_edit(
+        self, before: discord.Message, after: discord.Message
+    ) -> None:
+        await self.process_mentions(after)
 
 
 def main():