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 .

aocd/__init__.py

@@ -1,7 +1,7 @@
 import sys
 from functools import partial
+from typing import TYPE_CHECKING
 
-from . import _ipykernel
 from . import cli
 from . import cookies
 from . import examples
@@ -16,8 +16,28 @@
 from .get import get_day_and_year
 from .post import submit as _impartial_submit
 
+__all__ = [
+    "AocdError",
+    "cli",
+    "cookies",
+    "data",
+    "examples",
+    "exceptions",
+    "get_data",
+    "get",
+    "models",
+    "post",
+    "runner",
+    "submit",
+    "utils",
+]
 
-def __getattr__(name):
+data: str
+
+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/_compat.py

@@ -0,0 +1,40 @@
+import sys
+from importlib.metadata import entry_points
+
+# typing.Self added in 3.11
+if sys.version_info >= (3, 11):
+    from typing import Self as Self # import using same name to tell the type checker we intend to export this (so other modules can import it)
+else:
+    from typing_extensions import Self as Self # import using same name to tell the type checker we intend to export this (so other modules can import it)
+
+# typing.ParamSpec added in 3.10
+if sys.version_info >= (3, 10):
+    from typing import ParamSpec as ParamSpec # import using same name to tell the type checker we intend to export this (so other modules can import it)
+else:
+    from typing_extensions import ParamSpec as ParamSpec # import using same name to tell the type checker we intend to export this (so other modules can import it)
+
+# importlib.metadata.EntryPoints added in 3.10
+if sys.version_info >= (3, 10):
+    from importlib.metadata import EntryPoints
+
+    # Python 3.10+ - group/name selectable entry points
+    def get_entry_points(group: str, name: str) -> EntryPoints:
+        return entry_points().select(group=group, name=name)
+
+    def get_plugins(group: str = "adventofcode.user") -> EntryPoints:
+        """
+        Currently installed plugins for user solves.
+        """
+        return entry_points(group=group)
+else:
+    from importlib.metadata import EntryPoint
+
+    # Python 3.9 - dict interface
+    def get_entry_points(group: str, name: str) -> list[EntryPoint]:
+        return [ep for ep in entry_points()[group] if ep.name == name]
+
+    def get_plugins(group: str = "adventofcode.user") -> list[EntryPoint]:
+        """
+        Currently installed plugins for user solves.
+        """
+        return entry_points().get(group, [])

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,no-untyped-call] # 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")
         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,19 @@
+from decimal import Decimal
+from fractions import Fraction
+from typing import Literal
+from typing import Optional
+from typing import TYPE_CHECKING
+from typing import Union
+
+if TYPE_CHECKING:
+    import numpy as np
+
+__all__ = [
+    "_Answer",
+    "_Part",
+    "_LoosePart",
+]
+
+_Answer = Optional[Union[str, bytes, int, float, complex, Fraction, Decimal, "np.number[np.typing.NBitBase]"]]
+_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

@@ -5,7 +5,10 @@
 from dataclasses import dataclass
 from datetime import datetime
 from itertools import zip_longest
+from typing import Callable
 from typing import NamedTuple
+from typing import Optional
+from typing import Union
 
 import bs4
 
@@ -15,7 +18,6 @@
 from aocd.utils import AOC_TZ
 from aocd.utils import get_plugins
 
-
 log = logging.getLogger(__name__)
 
 
@@ -35,18 +37,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 +58,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[0]
+        assert isinstance(article_a, bs4.Tag)
+        a_raw = str(article_a)
+
+        article_b = b_raw = None
+        if len(articles) == 2:
+            article_b = articles[1]
+            assert isinstance(article_b, bs4.Tag)
+            b_raw = str(article_b)
+
         page = Page(
             raw_html=html,
             soup=soup,
@@ -77,7 +83,7 @@ def from_raw(cls, html):
         )
         return page
 
-    def __getattr__(self, name):
+    def __getattr__(self, name: str) -> Union[list[bs4.Tag], list[str]]:
         if not name.startswith(("a_", "b_")):
             raise AttributeError(name)
         part, sep, tag = name.partition("_")
@@ -91,13 +97,16 @@ 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
+        tags: list[bs4.Tag] = article.find_all(tag)
+        result: Union[list[bs4.Tag], list[str]] = tags
         if tag == "li":
             # list items usually need further drill-down
-            result = article.find_all("li")
-            for li in result:
-                li.codes = [code.text for code in li.find_all("code")]
+            for li in tags:
+                code_tags: list[bs4.Tag] = li.find_all("code")
+                li.codes = [code.text for code in code_tags] # type: ignore[attr-defined]
         else:
-            result = [t.text for t in article.find_all(tag)]
+            result = [t.text for t in tags]
         setattr(self, name, result)  # cache the result
         msg = "cached %s accessors for puzzle %d/%02d part %s page (%d hits)"
         log.debug(msg, tag, self.year, self.day, part, len(result))
@@ -118,23 +127,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 +153,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 +214,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)