Add type annotations

.github/workflows/tests.yml

@@ -34,10 +34,13 @@ jobs:
           architecture: x64
 
       - name: "Install"
-        run: pip install -q -r tests/requirements.txt && pip freeze --all
+        run: pip install -q -e .[all] && pip freeze --all
 
       - name: "Run tests for ${{ matrix.python-version }} on ${{ matrix.os }}"
         run: python -m pytest --durations=10
 
       - name: Upload coverage to Codecov
         uses: "codecov/codecov-action@main"
+
+      - name: "Run type check for ${{ matrix.python-version }} on ${{ matrix.os }}"
+        run: python -m mypy .

.gitignore

@@ -0,0 +1,4 @@
+*.egg-info
+__pycache__/
+.coverage*
+dist/

aocd/__init__.py

@@ -1,7 +1,7 @@
 import sys
+from typing import TYPE_CHECKING, Union
 from functools import partial
 
-from . import _ipykernel
 from . import cli
 from . import cookies
 from . import examples
@@ -16,8 +16,23 @@
 from .get import get_day_and_year
 from .post import submit as _impartial_submit
 
+__all__ = [
+    "exceptions",
+    "models",
+    "AocdError",
+    "get_data",
+    "data",
+    "submit",
+]
 
-def __getattr__(name):
+data: str
+lines: list[str]
+numbers: Union[list[list[int]], list[int], int]
+
+if TYPE_CHECKING:
+    submit = _impartial_submit
+
+def __getattr__(name): # type: ignore[no-untyped-def] # no idea how to provide meaningful types here
     if name == "data":
         day, year = get_day_and_year()
         return get_data(day=day, year=year)

aocd/_ipykernel.py

@@ -3,23 +3,31 @@
 
 import urllib3
 
+_IPYNB_PATTERN = re.compile(r"(?<=kernel-)[\w\-]+(?=\.json)")
 
-def get_ipynb_path():
+def get_ipynb_path() -> str:
     # helper function so that "from aocd import data" can introspect the year/day from
     # the .ipynb filename. inline imports to avoid hard-dependency on IPython/jupyter
     import IPython
     from jupyter_server import serverapp
     from jupyter_server.utils import url_path_join
 
-    app = IPython.get_ipython().config["IPKernelApp"]
-    kernel_id = re.search(r"(?<=kernel-)[\w\-]+(?=\.json)", app["connection_file"])[0]
+    ipython = IPython.get_ipython() # type: ignore[attr-defined] # IPython doesn't explicitly export the get_ipython function
+    app = ipython.config["IPKernelApp"]
+    match = _IPYNB_PATTERN.search(app["connection_file"])
+    assert match is not None
+    kernel_id = match[0]
     http = urllib3.PoolManager()
-    for serv in serverapp.list_running_servers():
-        url = url_path_join(serv["url"], "api/sessions")
+    for serv in serverapp.list_running_servers(): # type: ignore[no-untyped-call] # jupyter doesn't provide type hints for list_running_servers
+        url = url_path_join(serv["url"], "api/sessions") # type: ignore[no-untyped-call] # jupyter doesn't provide type hints for url_path_join
         resp = http.request("GET", url, fields={"token": serv["token"]})
-        resp.raise_for_status()
+        if resp.status >= 400:
+            raise urllib3.exceptions.ResponseError(f"Bad HTTP response status ({resp.status})")
         for sess in resp.json():
             if kernel_id == sess["kernel"]["id"]:
                 path = serv["root_dir"]
+                assert isinstance(path, str)
                 fname = sess["notebook"]["path"]
+                assert isinstance(fname, str)
                 return os.path.join(path, fname)
+    assert False, "unreachable"

aocd/_types.py

@@ -0,0 +1,14 @@
+from typing import TYPE_CHECKING, Literal, Optional, Union
+
+if TYPE_CHECKING:
+    import numpy as np
+
+__all__ = [
+    "_Answer",
+    "_Part",
+    "_LoosePart",
+]
+
+_Answer = Optional[Union[str, int, float, complex, "np.number"]]
+_Part = Literal["a", "b"]
+_LoosePart = Union[_Part, Literal[1, "1", 2, "2"]]

aocd/cli.py

@@ -13,7 +13,7 @@
 from .utils import get_plugins
 
 
-def main():
+def main() -> None:
     """Get your puzzle input data, caching it if necessary, and print it on stdout."""
     aoc_now = datetime.datetime.now(tz=AOC_TZ)
     days = range(1, 26)

aocd/cookies.py

@@ -4,18 +4,21 @@
 import logging
 import os
 import sys
+from typing import TYPE_CHECKING
 
 from .exceptions import DeadTokenError
 from .models import AOCD_CONFIG_DIR
 from .utils import _ensure_intermediate_dirs
 from .utils import colored
 from .utils import get_owner
 
+if TYPE_CHECKING:
+    import http.cookiejar
 
 log = logging.getLogger(__name__)
 
 
-def get_working_tokens():
+def get_working_tokens() -> dict[str, str]:
     """Check browser cookie storage for session tokens from .adventofcode.com domain."""
     log.debug("checking for installation of browser-cookie3 package")
     try:
@@ -25,9 +28,8 @@ def get_working_tokens():
 
     log.info("checking browser storage for tokens, this might pop up an auth dialog!")
     log.info("checking chrome cookie jar...")
-    cookie_files = glob.glob(os.path.expanduser("~/.config/google-chrome/*/Cookies"))
-    cookie_files.append(None)
-    chrome_cookies = []
+    cookie_files = (*glob.glob(os.path.expanduser("~/.config/google-chrome/*/Cookies")), None)
+    chrome_cookies: list["http.cookiejar.Cookie"] = []
     for cf in cookie_files:
         try:
             chrome = bc3.chrome(cookie_file=cf, domain_name=".adventofcode.com")
@@ -49,12 +51,12 @@ def get_working_tokens():
         log.info("%d candidates from firefox", len(firefox))
 
     # order preserving de-dupe
-    tokens = list({}.fromkeys([c.value for c in chrome + firefox]))
+    tokens = list({}.fromkeys([c.value for c in chrome + firefox if c.value is not None]))
     removed = len(chrome + firefox) - len(tokens)
     if removed:
         log.info("Removed %d duplicate%s", removed, "s"[: removed - 1])
 
-    result = {}  # map of {token: auth source}
+    result: dict[str, str] = {}  # map of {token: auth source}
     for token in tokens:
         try:
             owner = get_owner(token)
@@ -66,7 +68,7 @@ def get_working_tokens():
     return result
 
 
-def scrape_session_tokens():
+def scrape_session_tokens() -> None:
     """Scrape AoC session tokens from your browser's cookie storage."""
     aocd_token_path = AOCD_CONFIG_DIR / "token"
     aocd_tokens_path = AOCD_CONFIG_DIR / "tokens.json"
@@ -104,7 +106,7 @@ def scrape_session_tokens():
             if aocd_token_path.is_file():
                 txt = aocd_token_path.read_text(encoding="utf-8").strip()
                 if txt:
-                    tokens[aocd_token_path] = txt.split()[0]
+                    tokens[str(aocd_token_path)] = txt.split()[0]
             if aocd_tokens_path.is_file():
                 tokens.update(json.loads(aocd_tokens_path.read_text(encoding="utf-8")))
         else:

aocd/examples.py

@@ -2,10 +2,10 @@
 import logging
 import re
 import sys
+from typing import NamedTuple, Optional, Callable
 from dataclasses import dataclass
 from datetime import datetime
 from itertools import zip_longest
-from typing import NamedTuple
 
 import bs4
 
@@ -15,7 +15,6 @@
 from aocd.utils import AOC_TZ
 from aocd.utils import get_plugins
 
-
 log = logging.getLogger(__name__)
 
 
@@ -35,18 +34,19 @@ class Page:
     year: int  # AoC puzzle year (2015+) parsed from html title
     day: int  # AoC puzzle day (1-25) parsed from html title
     article_a: bs4.element.Tag  # The bs4 tag for the first <article> in the page, i.e. part a
-    article_b: bs4.element.Tag  # The bs4 tag for the second <article> in the page, i.e. part b. It will be `None` if part b locked
+    article_b: Optional[bs4.element.Tag]  # The bs4 tag for the second <article> in the page, i.e. part b. It will be `None` if part b locked
     a_raw: str  # The first <article> html as a string
-    b_raw: str  # The second <article> html as a string. Will be `None` if part b locked
+    b_raw: Optional[str]  # The second <article> html as a string. Will be `None` if part b locked
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         part_a_only = "*" if self.article_b is None else ""
         return f"<Page({self.year}, {self.day}){part_a_only} at {hex(id(self))}>"
 
     @classmethod
-    def from_raw(cls, html):
+    def from_raw(cls, html: str) -> "Page":
         soup = _get_soup(html)
         title_pat = r"^Day (\d{1,2}) - Advent of Code (\d{4})$"
+        assert soup.title is not None
         title_text = soup.title.text
         if (match := re.match(title_pat, title_text)) is None:
             msg = f"failed to extract year/day from title {title_text!r}"
@@ -55,16 +55,19 @@ def from_raw(cls, html):
         articles = soup.find_all("article")
         if len(articles) == 0:
             raise ExampleParserError("no <article> found in html")
-        elif len(articles) == 1:
-            [article_a] = articles
-            a_raw = str(article_a)
-            article_b = b_raw = None
-        elif len(articles) == 2:
-            article_a, article_b = articles
-            a_raw = str(article_a)
-            b_raw = str(article_b)
-        else:
+        if len(articles) > 2:
             raise ExampleParserError("too many <article> found in html")
+
+        article_a = articles.pop(0)
+        assert isinstance(article_a, bs4.Tag)
+        a_raw = str(article_a)
+
+        article_b = b_raw = None
+        if articles:
+            article_b = articles.pop(0)
+            assert isinstance(article_b, bs4.Tag)
+            b_raw = str(article_b)
+
         page = Page(
             raw_html=html,
             soup=soup,
@@ -77,7 +80,7 @@ def from_raw(cls, html):
         )
         return page
 
-    def __getattr__(self, name):
+    def __getattr__(self, name: str) -> list[bs4.Tag]:
         if not name.startswith(("a_", "b_")):
             raise AttributeError(name)
         part, sep, tag = name.partition("_")
@@ -91,11 +94,13 @@ def __getattr__(self, name):
             # actually used by an example parser
             raise AttributeError(name)
         article = self.article_a if part == "a" else self.article_b
+        assert article is not None
         if tag == "li":
             # list items usually need further drill-down
-            result = article.find_all("li")
+            result: list[bs4.Tag] = article.find_all("li")
             for li in result:
-                li.codes = [code.text for code in li.find_all("code")]
+                codes: bs4.ResultSet[bs4.Tag] = li.find_all("code")
+                li.codes = [code.text for code in codes] # type: ignore[attr-defined]
         else:
             result = [t.text for t in article.find_all(tag)]
         setattr(self, name, result)  # cache the result
@@ -118,23 +123,23 @@ class Example(NamedTuple):
     """
 
     input_data: str
-    answer_a: str = None
-    answer_b: str = None
-    extra: str = None
+    answer_a: Optional[str] = None
+    answer_b: Optional[str] = None
+    extra: Optional[str] = None
 
     @property
-    def answers(self):
+    def answers(self) -> tuple[Optional[str], Optional[str]]:
         return self.answer_a, self.answer_b
 
 
-def _trunc(s, maxlen=50):
+def _trunc(s: Optional[str], maxlen: int = 50) -> Optional[str]:
     # don't print massive strings and mess up the table rendering
     if s is None or len(s) <= maxlen:
         return s
     return s[:maxlen] + f" ... ({len(s)} bytes)"
 
 
-def _get_unique_real_inputs(year, day):
+def _get_unique_real_inputs(year: int, day: int) -> list[str]:
     # these are passed to example parsers, in case the shape/content of the real
     # input(s) is in some way useful for extracting the example input(s). it is
     # not currently used by the default example parser implementation.
@@ -144,7 +149,7 @@ def _get_unique_real_inputs(year, day):
     return list({}.fromkeys(strs))
 
 
-def main():
+def main() -> None:
     """
     Summarize an example parser's results with historical puzzles' prose, and
     compare the performance against a reference implementation
@@ -205,7 +210,7 @@ def main():
             file=sys.stderr,
         )
         sys.exit(1)
-    plugin = plugins[args.example_parser].load()
+    plugin: Callable[[Page, list[str]], list[Example]] = plugins[args.example_parser].load()
     console = Console()
     parser_wants_real_datas = getattr(plugin, "uses_real_datas", True)