add type hints to tests

Author: afeld

extras/lib/nb_helper.py

@@ -1,41 +1,42 @@
 import nbconvert
 import nbformat
+from nbformat import NotebookNode
 from .diffable import is_system_command
 
 
-def notebook_to_script(notebook):
+def notebook_to_script(notebook: NotebookNode):
     exporter = nbconvert.exporters.PythonExporter
     script, _ = nbconvert.exporters.export(exporter, notebook)
     return script
 
 
-def read_notebook(notebook_path):
+def read_notebook(notebook_path: str) -> NotebookNode:
     return nbformat.read(notebook_path, as_version=4)
 
 
-def is_markdown(cell):
+def is_markdown(cell: NotebookNode):
     return cell.cell_type == "markdown"
 
 
-def is_h1(cell):
+def is_h1(cell: NotebookNode):
     return is_markdown(cell) and cell.source.startswith("# ")
 
 
-def is_magic(source):
+def is_magic(source: str):
     """https://ipython.readthedocs.io/en/stable/interactive/tutorial.html#magic-functions"""
     return source.startswith("%%")
 
 
-def is_code_cell(cell):
+def is_code_cell(cell: NotebookNode):
     return cell.cell_type == "code"
 
 
-def is_python(cell):
+def is_python(cell: NotebookNode):
     source = cell.source
     return cell.cell_type == "code" and not (
         is_magic(source) or is_system_command(source)
     )
 
 
-def get_tags(cell):
+def get_tags(cell: NotebookNode):
     return cell.metadata.get("tags", [])

extras/test_lectures.py

@@ -1,6 +1,7 @@
 from glob import glob
 import re
 
+from nbformat import NotebookNode
 import pytest
 
 from .lib.nb_helper import get_tags, is_code_cell, read_notebook
@@ -10,11 +11,11 @@
 lecture_notebooks.sort()
 
 
-def slide_type(cell):
+def slide_type(cell: NotebookNode):
     return cell.metadata.get("slideshow", {}).get("slide_type")
 
 
-def is_slide(cell):
+def is_slide(cell: NotebookNode):
     SLIDE_TYPES = ["slide", "subslide"]
     return slide_type(cell) in SLIDE_TYPES
 
@@ -23,7 +24,7 @@ def is_exercise(source: str):
     return re.match("#+.+(demo|exercise)", source, re.IGNORECASE) is not None
 
 
-def num_slides(cells):
+def num_slides(cells: list[NotebookNode]):
     """Return a weighted number of slides"""
 
     slides = [cell for cell in cells if is_slide(cell)]
@@ -51,12 +52,13 @@ def test_num_slides(file):
     if file == "lecture_6.ipynb":
         pytest.xfail("The various pieces of the lecture can be scaled appropriately")
 
-    num_columbia = num_slides_without_tag(notebook.cells, "nyu-only")
+    cells: list[NotebookNode] = notebook.cells
+    num_columbia = num_slides_without_tag(cells, "nyu-only")
     print("Number of slides for Columbia: ", num_columbia)
     assert num_columbia >= 42, "Too few slides for Columbia"
     assert num_columbia <= 63, "Too many slides for Columbia"
 
-    num_nyu = num_slides_without_tag(notebook.cells, "columbia-only")
+    num_nyu = num_slides_without_tag(cells, "columbia-only")
     print("Number of slides for NYU: ", num_nyu)
     assert num_nyu >= 39, "Too few slides for NYU"
     assert num_nyu <= 51, "Too many slides for NYU"

pyproject.toml

@@ -1,3 +1,7 @@
+# https://mypy.readthedocs.io/en/stable/config_file.html
+[tool.mypy]
+check_untyped_defs = true
+
 [tool.pytest.ini_options]
 addopts = [
     "--ignore=extras/autograder/",