Collect Notebooks and Cells with pytest (#7)

Author: MusicalNinjaDad

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added  
 
+- Pytest plugin can collect jupyter notebooks and cells which use the `%%ipytest` magic
 - Published documentation for `CollectionTree`.
 
 ## [0.0.2] - 2025-02-12

docs/parser.md

@@ -0,0 +1,5 @@
+# Parser (internal)
+
+The parser will likely be deprecated and merged into the plug-in classes in the future
+
+::: pytest_ipynb2._parser

docs/plugin.md

@@ -0,0 +1,3 @@
+# Pytest plugin
+
+::: pytest_ipynb2.plugin

docs/python-api.md

@@ -48,6 +48,10 @@ check:
   @- just test
   @- just type-check
 
+# format with ruff
+format:
+  uv run ruff format .
+
 #run coverage analysis on python code
 cov:
   uv run pytest --cov --cov-report html:pycov --cov-report term --cov-context=test

justfile

@@ -7,8 +7,9 @@ watch:
 
 nav:
   - index.md
-  - python-api.md
+  - plugin.md
   - test_helpers.md
+  - parser.md
 
 theme:
   name: "material"

mkdocs.yml

@@ -81,6 +81,7 @@
         "--doctest-mdcodeblocks",
         "--doctest-glob=*.pyi",
         "--doctest-glob=*.md",
+        "--ignore=tests/assets",
     ]
     testpaths = ["tests", "pytest_ipynb2"]
 

pyproject.toml

@@ -1,3 +1 @@
 """Pytest plugin to run tests in Jupyter Notebooks."""
-
-from ._ipynb_parser import Notebook

pytest_ipynb2/__init__.py

@@ -6,10 +6,13 @@
 import nbformat
 
 
+class Cell: ...
+
+
 class Notebook:
     """
     An ipython Notebook.
-     
+
     - constructor from Path
     - methods to get various cell types
     - a `test` cell type identified by the `%%ipytest` cell magic.
@@ -35,4 +38,4 @@ def gettestcells(self) -> dict[int, str]:
             cellnr: "".join(cell["source"][1:]).strip()
             for cellnr, cell in enumerate(self.contents["cells"])
             if cell["cell_type"] == "code" and cell["source"][0].startswith(r"%%ipytest")
-        }
+        }

pytest_ipynb2/_ipynb_parser.py renamed to pytest_ipynb2/_parser.py

@@ -0,0 +1,64 @@
+"""
+Pytest plugin to collect jupyter Notebooks.
+
+- Identifies all cells which use the `%%ipytest` magic
+- adds the notebook, cell and any test functions to the collection tree
+- relies on pytest logic and configuration to identify test functions.
+"""
+
+from __future__ import annotations
+
+import importlib
+from typing import TYPE_CHECKING
+
+import pytest
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+    from pathlib import Path
+    from types import ModuleType
+
+from ._parser import Notebook as _ParsedNotebook
+
+ipynb2_notebook = pytest.StashKey[_ParsedNotebook]()
+
+
+class Notebook(pytest.File):
+    """A collector for jupyter notebooks."""
+
+    def collect(self) -> Generator[Cell, None, None]:
+        """Yield `Cell`s for all cells which contain tests."""
+        parsed = _ParsedNotebook(self.path)
+        for testcellid in parsed.gettestcells():
+            cell = Cell.from_parent(
+                parent=self,
+                name=str(testcellid),
+                nodeid=str(testcellid),
+                path=self.path,
+            )
+            cell.stash[ipynb2_notebook] = parsed
+            yield cell
+
+
+class Cell(pytest.Module):
+    """
+    A collector for jupyter notebook cells.
+    
+    `pytest` will recognise these cells as `pytest.Module`s and use standard collection on them as it would any other
+    python module.
+    """
+
+    def _getobj(self) -> ModuleType:
+        notebook = self.stash[ipynb2_notebook]
+        cellsource = notebook.gettestcells()[int(self.nodeid)]
+        cellspec = importlib.util.spec_from_loader(f"Cell{self.name}", loader=None)
+        cell = importlib.util.module_from_spec(cellspec)
+        exec(cellsource, cell.__dict__)  # noqa: S102
+        return cell
+
+
+def pytest_collect_file(file_path: Path, parent: pytest.Collector) -> Notebook | None:
+    """Hook implementation to collect jupyter notebooks."""
+    if file_path.suffix == ".ipynb":
+        return Notebook.from_parent(parent=parent, path=file_path)
+    return None

pytest_ipynb2/plugin.py

@@ -10,10 +10,12 @@
 
 if TYPE_CHECKING:
     from contextlib import suppress
+    from pathlib import Path
 
-    with suppress(ImportError): # not type-checking on python < 3.11
+    with suppress(ImportError):  # not type-checking on python < 3.11
         from typing import Self
 
+
 class CollectionTree:
     """
     A (top-down) tree of pytest collection Nodes.
@@ -25,25 +27,29 @@ class CollectionTree:
     """
 
     @classmethod
+    #TODO(MusicalNinjaDad): #8 Refactor CollectionTree.from_items to be easier to understand.
     def from_items(cls, items: list[pytest.Item]) -> Self:
         """Create a CollectionTree from a list of collection items, as returned by `pytester.genitems()`."""
+
         def _from_item(item: pytest.Item | Self) -> Self:
             return item if isinstance(item, cls) else cls(node=item, children=None)
-        
-        items = [_from_item(item) for item in items]
 
         def _get_parents_as_CollectionTrees(items: list[Self]) -> list[Self]:  # noqa: N802
             """Returns a list of CollectionTree items for the parents of those provided."""
             parents = (item.node.parent for item in items)
             items_byparent = {
                 parent: [item for item in items if item.node.parent == parent]
                 for parent in parents
-            }
-            return [
-                cls(node = parent, children = list(children))
-                for parent, children in items_byparent.items()
-            ]
-        
+            }  # fmt: skip
+            return [cls(node=parent, children=list(children)) for parent, children in items_byparent.items()]
+
+        if not items:
+            # If we don't specifically handle this here, then `all([])` returns `True` later and wierd stuff happens
+            msg = "Items list is empty."
+            raise ValueError(msg)
+
+        items = [_from_item(item) for item in items]
+
         if all(isinstance(item.node, pytest.Session) for item in items):
             assert len(items) == 1, "There should only ever be one session node."  # noqa: S101
             return next(iter(items))
@@ -93,7 +99,7 @@ def from_dict(cls, tree: dict[tuple[str, type], dict | None]) -> Self:
                     cls.from_dict({childnode: grandchildren})
                     for childnode, grandchildren in children.items()
                 ],
-            )
+            )  # fmt:skip
 
         return cls(node=node, children=None)
 
@@ -155,3 +161,26 @@ def __repr__(self) -> str:
         else:
             children_repr = indent("\n".join(repr(child).rstrip() for child in self.children), "    ")
         return f"{self.node!r}\n{children_repr}\n"
+
+
+@dataclass
+class CollectedDir:
+    """
+    The various elements required to test directory collection.
+
+    - `pytester_instance`: pytest.Pytester
+    - `dir_node`: pytest.Dir
+    - `items`: list[pytest.Item]
+    """
+
+    pytester_instance: pytest.Pytester
+    dir_node: pytest.Dir
+    items: list[pytest.Item]
+
+
+@dataclass
+class ExampleDir:
+    """The various elements to set up a pytester instance."""
+
+    files: list[Path]
+    conftest: str = ""

pytest_ipynb2/pytester_helpers.py

@@ -18,7 +18,7 @@
     "x = 1\n",
     "y = 2\n",
     "\n",
-    "x + y\n"
+    "x + y"
    ]
   },
   {
@@ -36,8 +36,9 @@
    "source": [
     "# Define a function\n",
     "\n",
+    "\n",
     "def adder(a, b):\n",
-    "    return a+b"
+    "    return a + b"
    ]
   },
   {
@@ -48,8 +49,10 @@
    "source": [
     "%%ipytest\n",
     "\n",
+    "\n",
     "def test_adder():\n",
-    "    assert adder(1,2) == 3\n",
+    "    assert adder(1, 2) == 3\n",
+    "\n",
     "\n",
     "def test_globals():\n",
     "    assert x == 1"
@@ -82,7 +85,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.1"
+   "version": "3.13.2"
   }
  },
  "nbformat": 4,

tests/assets/notebook.ipynb

@@ -1,4 +1,4 @@
-#ruff: noqa
+# ruff: noqa
 # This cell sets some global variables
 
 x = 1
@@ -8,14 +8,18 @@
 
 # Define a function
 
+
 def adder(a, b):
-    return a+b
+    return a + b
+
 
 def test_adder():
-    assert adder(1,2) == 3
+    assert adder(1, 2) == 3
+
 
 def test_globals():
     assert x == 1
 
+
 def another_function(*args):
-    return args
+    return args