Examples page

.gitignore

@@ -1,2 +1,12 @@
 # Files generated by Sphinx
 build/
+source/_static/examples/
+source/examples/
+
+# Python stuff
+__pycache__/
+.pytest_cache
+.ipynb_checkpoints
+
+# SOAP
+.soap/

devtools/conda-envs/examples_env.yml

@@ -0,0 +1,16 @@
+channels:
+    - conda-forge
+    - bioconda
+dependencies:
+    - pip
+    - python=3.10
+    # Cookbook
+    - gitpython
+    - nbconvert
+    - nbformat
+    # Examples
+    - openff-toolkit-examples
+    - gromacs
+    - lammps
+    - rich
+    - jax

devtools/conda-envs/lint_env.yml

@@ -7,6 +7,8 @@ dependencies:
     - sphinx>=5,<7
     - myst-parser>=1,<2
     - sphinx-notfound-page
+    - nbsphinx
+    - gitpython
     # Code example deps
     - openff-toolkit-base
     - openff-interchange-base

devtools/conda-envs/rtd_env.yml

@@ -7,6 +7,12 @@ dependencies:
     - sphinx>=5,<7
     - myst-parser>=1,<2
     - sphinx-notfound-page
+    - nbsphinx
+    - ipython >=8.8
+    # Examples
+    - gitpython
+    - nbconvert
+    - nbformat
     # Theme
     - pip:
         # Theme

source/_ext/cookbook/__init__.py

@@ -0,0 +1,60 @@
+from ._cookbook import (
+    collect_notebooks,
+    remove_old_notebooks,
+    process_notebook,
+    find_notebook_docnames,
+    CookbookDirective,
+)
+from sphinx.application import Sphinx as Application
+
+
+def setup(app: Application):
+    # TODO: Remove these config values and replace with simpler versions
+    app.add_config_value(
+        "cookbook_default_conda_forge_deps",
+        default=[],
+        rebuild="env",
+    )
+    app.add_config_value(
+        "cookbook_default_required_files",
+        default=None,
+        rebuild="env",
+    )
+    app.add_config_value(
+        "cookbook_required_files_base_uri",
+        default="",
+        rebuild="env",
+    )
+
+    app.add_config_value(
+        "cookbook_example_repos",
+        default=[],
+        rebuild="env",
+    )
+    app.add_config_value(
+        "cookbook_examples_path",
+        default="_static/examples/",
+        rebuild="env",
+    )
+    app.add_config_value(
+        "cookbook_dont_fetch",
+        default=False,
+        rebuild="env",
+    )
+    # TODO: Implement this config value
+    app.add_config_value(
+        "cookbook_ignore_notebook_pattern",
+        default=["deprecated/**"],
+        rebuild="env",
+    )
+
+    app.connect("config-inited", collect_notebooks)
+    app.connect("env-purge-doc", remove_old_notebooks)
+    app.connect("source-read", process_notebook)
+    app.connect("env-before-read-docs", find_notebook_docnames)
+    app.add_directive("cookbook", CookbookDirective)
+
+    return {
+        "parallel_read_safe": True,
+        "parallel_write_safe": True,
+    }

source/_ext/cookbook/_cookbook.py

@@ -0,0 +1,113 @@
+"""Implementation of the cookbook Sphinx extension"""
+
+import json
+from pathlib import Path
+from typing import Generator
+
+from sphinx.application import Sphinx as Application
+from sphinx.config import Config
+from sphinx.environment import BuildEnvironment
+from nbsphinx import NbGallery
+
+from ._notebook import insert_cell, get_metadata, find_notebooks, notebook_zip
+from ._github import download_dir
+
+
+def inject_tags_index(notebook: dict) -> dict:
+    """Inject an `index` directive containing the notebook's metadata tags"""
+
+    tags = get_metadata(notebook, "tags", ["untagged"])
+
+    return insert_cell(
+        notebook,
+        cell_type="raw",
+        metadata={"raw_mimetype": "text/restructuredtext"},
+        source=[
+            f".. index:: {', '.join(tags)}",
+        ],
+    )
+
+
+def inject_links(notebook: dict, docpath: Path, cookbook_examples_path: Path) -> dict:
+    user, repo, *path = str(docpath.relative_to(cookbook_examples_path)).split("/")
+    path = "/".join(path)
+
+    github_url = f"https://github.com/{user}/{repo}/blob/main/{path}"
+    # TODO: Figure out how to get the conda colab install cell into this
+    colab_url = (
+        f"https://colab.research.google.com/github/{user}/{repo}/blob/main/{path}"
+    )
+
+    return insert_cell(
+        notebook,
+        cell_type="raw",
+        metadata={"raw_mimetype": "text/restructuredtext"},
+        source=[
+            f":download:`Download Notebook </{notebook_zip(docpath)}>`",
+            f"`View in GitHub <{github_url}>`_",
+            f"`Open in Google Colab <{colab_url}>`_",
+        ],
+    )
+
+
+def process_notebook(app: Application, docname: str, source: list[str]):
+    docpath = Path(app.env.doc2path(docname, False))
+    if docpath.suffix != ".ipynb":
+        return
+
+    notebook = json.loads(source[0])
+
+    notebook = inject_links(notebook, docpath, app.config.cookbook_examples_path)
+    notebook = inject_tags_index(notebook)
+
+    source[0] = json.dumps(notebook)
+
+
+def remove_old_notebooks(app: Application, env: BuildEnvironment, docname: str):
+    """Clean up processed notebooks from outdir"""
+    # outdir = Path(app.outdir)
+    # colab_path = outdir / "colab" / docname
+    # colab_path = colab_path.with_suffix(".ipynb")
+    # # We check is_relative_to just in case docname is absolute
+    # if colab_path.exists() and colab_path.is_relative_to(outdir / "colab"):
+    #     colab_path.unlink()
+
+
+def collect_notebooks(app: Application, config: Config):
+    """Download examples directories from example repositories"""
+    app.cookbook_examples_path = Path(app.srcdir) / config.cookbook_examples_path  # type: ignore
+    if config.cookbook_dont_fetch:
+        return
+    for repo in config.cookbook_example_repos:
+        print("Collecting examples from", repo)
+
+        download_dir(repo, "examples", app.cookbook_examples_path)  # type: ignore
+
+
+def find_notebook_docnames(app, env, docnames):
+    """Find the downloaded notebooks and make sure Sphinx sees them"""
+    if not hasattr(app, "cookbook_notebooks"):
+        app.cookbook_notebooks = []  # type: ignore
+
+    for repo in env.config.cookbook_example_repos:
+        for path in find_notebooks(app.cookbook_examples_path / repo / "examples"):
+            docname = env.project.path2doc(str(path.relative_to(app.srcdir)))
+            app.cookbook_notebooks.append(docname)
+            docnames.append(docname)
+
+
+class CookbookDirective(NbGallery):
+    """
+    Directive to draw thumbnails of the cookbook.
+    """
+
+    def run(self):
+        nodes = super().run()
+        toctree = nodes[0][0][0]
+
+        toctree["entries"].extend(
+            [(None, docname) for docname in self.env.app.cookbook_notebooks]
+        )
+        toctree["includefiles"].extend(self.env.app.cookbook_notebooks)
+
+        return nodes

source/_ext/cookbook/_github.py

@@ -0,0 +1,124 @@
+"""Code for working with GitHub in both the Sphinx extension and proc_examples.py"""
+import shutil
+from pathlib import Path
+from typing import Generator, Iterable, Literal, Optional
+
+from git.repo import Repo
+from git.diff import DiffIndex
+
+
+class UpdateSpec:
+    """Stores paths that must be reprocessed or cleaned up"""
+
+    def __init__(
+        self,
+        reprocess: Optional[Iterable[Path]] = None,
+        cleanup: Optional[Iterable[Path]] = None,
+        start_over=False,
+    ):
+        # Find the notebooks associated with each change, and remove duplicates
+        reprocess = set(self._find_all_notebooks(reprocess))
+        cleanup = set(self._find_all_notebooks(cleanup))
+
+        # We don't need to cleanup notebooks we're reprocessing, so remove them
+        cleanup = cleanup - reprocess
+
+        self.reprocess = list(reprocess)
+        """List of paths to reprocess"""
+        self.cleanup = list(cleanup)
+        """List of paths to clean up"""
+        self.start_over = start_over
+        """If True, all notebooks should be cleaned up and reprocessed"""
+
+    @classmethod
+    def from_diff(cls, diffs) -> "UpdateSpec":
+        reprocess = []
+        cleanup = []
+
+        for diff in diffs:
+            # Files we need to re-process (from additions, copies, renames,
+            # modifications)
+            if diff.change_type in "ACRM":
+                reprocess.append(Path(diff.b_path))
+            # Files we need to clean up (from deletions, renames)
+            if diff.change_type == "DR":
+                cleanup.append(Path(diff.a_path))
+
+        return cls(reprocess=reprocess, cleanup=cleanup)
+
+    @classmethod
+    def update_all(cls) -> "UpdateSpec":
+        return cls(start_over=True)
+
+    @classmethod
+    def _find_all_notebooks(
+        cls, paths: Optional[Iterable[Path]]
+    ) -> Generator[Path, None, None]:
+        """Find the notebooks that could correspond to each file"""
+        if paths is None:
+            return
+
+        for path in paths:
+            yield from cls._find_notebooks(path)
+
+    @staticmethod
+    def _find_notebooks(path: Path) -> list[Path]:
+        """Find the notebooks that could correspond to a file"""
+        if path.suffix == ".ipynb":
+            return [path]
+
+        for parent in path.parents:
+            notebooks = [*parent.glob("*.ipynb")]
+            if notebooks:
+                return notebooks
+        return []
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}(reprocess={list(self.reprocess)},"
+            + f" cleanup={list(self.cleanup)}, start_over={self.start_over})"
+        )
+
+
+def download_dir(src_repo: str, src_path: str, examples_path: Path) -> UpdateSpec:
+    """Download src_path from GitHub src_repo
+
+    The folder is downloaded to to ``<examples_path>/<src_repo>/<src_path>``.
+    Git is used to only download new or updated files. The names of
+    newly-downloaded or updated files are returned."""
+    local_repo_path = examples_path / src_repo
+
+    # Try a simple pull
+    try:
+        repo = Repo(local_repo_path)
+
+        previous = repo.head.commit
+        repo.remotes.origin.pull(depth=1)
+        new = repo.head.commit
+    except Exception:
+        # Clean up whatever's there
+        if local_repo_path.exists():
+            shutil.rmtree(local_repo_path)
+    else:
+        if previous == new:
+            return UpdateSpec()
+        else:
+            return UpdateSpec.from_diff(previous.diff(new))
+
+    # If the repo doesn't exist, or a pull fails for some other reason, clone
+    local_repo_path.mkdir(parents=True, exist_ok=True)
+    # Clone without downloading anything
+    repo = Repo.clone_from(
+        url=f"https://github.com/{src_repo}.git",
+        to_path=local_repo_path,
+        multi_options=[
+            "--depth=1",
+            "--filter=tree:0",
+            "--no-checkout",
+        ],
+    )
+    # Just checkout the examples directory
+    repo.git.sparse_checkout("set", src_path)
+    repo.git.checkout()
+
+    return UpdateSpec.update_all()