Begin improvement of refs generation (#59)

* Begin improvement of refs generation

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Proof of concept works! Converts backtics to refs

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Simplify, trim stuff

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Add a test for the autoref plugin

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Remove unused code paths

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Update pre-commit config?

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Try a fix for broken docformatter pre-commit hook

- PyCQA/docformatter#289
- PyCQA/docformatter#287

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Restrict pre-commit version for now

Trying the same approach as DSD-DBS/capella-dockerimages@6e38a0c

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Temporarily don't use existing regression files

There's this issue at the moment with some reproducibility test
that cause the CI to fail.

I'll make another PR to address this, but for now, I'm turning off
the 'shared regression files in $SCRATCH' feature.

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Add test to reach 100% patch coverage

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

---------

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

.github/workflows/build.yml

@@ -27,7 +27,7 @@ jobs:
       - uses: actions/setup-python@v4
         with:
           python-version: "3.10"
-      - run: pip install pre-commit
+      - run: pip install 'pre-commit<4.0.0'
       - run: pre-commit --version
       - run: pre-commit install
       - run: pre-commit run --all-files

.pre-commit-config.yaml

@@ -3,7 +3,7 @@ default_language_version:
 
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
+    rev: v5.0.0
     hooks:
       # list of supported hooks: https://pre-commit.com/hooks.html
       - id: trailing-whitespace
@@ -33,7 +33,7 @@ repos:
 
   - repo: https://github.com/charliermarsh/ruff-pre-commit
     # Ruff version.
-    rev: "v0.5.7"
+    rev: "v0.6.9"
     hooks:
       - id: ruff
         args: ["--line-length", "99", "--fix"]
@@ -44,6 +44,7 @@ repos:
     rev: v1.7.5
     hooks:
       - id: docformatter
+        language: python
         args: [--in-place, --wrap-summaries=99, --wrap-descriptions=99]
         require_serial: true
 

mkdocs.yml

@@ -64,6 +64,7 @@ plugins:
   - awesome-pages
   - macros: #https://mkdocs-macros-plugin.readthedocs.io/en/latest/#declaration-of-the-macros-plugin
       module_name: docs/macros
+  - custom_autoref_plugin
   - autorefs
   - gen-files:
       # https://oprypin.github.io/mkdocs-gen-files/#usage

project/conftest.py

@@ -123,7 +123,10 @@ def original_datadir(original_datadir: Path):
     """
     relative_portion = original_datadir.relative_to(REPO_ROOTDIR)
     datadir = REPO_ROOTDIR / ".regression_files"
-    if SCRATCH and not datadir.exists():
+
+    use_scratch: bool = False  # todo: enable again later?
+
+    if SCRATCH and not datadir.exists() and use_scratch:
         # puts a symlink .regression_files in the repo root that points to the same dir in $SCRATCH
         actual_dir = SCRATCH / datadir.relative_to(REPO_ROOTDIR)
         actual_dir.mkdir(parents=True, exist_ok=True)

project/utils/autoref_plugin.py

@@ -0,0 +1,91 @@
+"""IDEA: Tweak the AutoRefsPlugin so that text in backticks like `this` (more IDE-friendly) are
+considered refs when possible.
+"""
+
+import functools
+import re
+
+import lightning
+import torch
+from mkdocs.config.defaults import MkDocsConfig
+from mkdocs.plugins import (
+    BasePlugin,
+    get_plugin_logger,
+)
+from mkdocs.structure.files import Files
+from mkdocs.structure.pages import Page
+from mkdocs_autorefs.plugin import AutorefsPlugin  # noqa
+
+from project.utils.hydra_config_utils import import_object
+
+# Same as in the mkdocs_autorefs plugin.
+logger = get_plugin_logger(__name__)
+
+known_things = [
+    lightning.Trainer,
+    lightning.LightningModule,
+    lightning.LightningDataModule,
+    torch.nn.Module,
+]
+"""
+IDEA: IF we see `Trainer`, and we know that that's the `lightning.Trainer`, then we
+create the proper ref.
+
+TODO: Ideally this would contain every object / class that we know of in this project.
+"""
+
+
+class CustomAutoRefPlugin(BasePlugin):
+    """Small mkdocs plugin that converts backticks to refs when possible."""
+
+    def on_page_markdown(
+        self, markdown: str, /, *, page: Page, config: MkDocsConfig, files: Files
+    ) -> str | None:
+        # Find all instances where backticks are used and try to convert them to refs.
+
+        # Examples:
+        # - `package.foo.bar` -> [package.foo.bar][] (only if `package.foo.bar` is importable)
+        # - `baz` -> [baz][]
+
+        def _full_path(thing) -> str:
+            return thing.__module__ + "." + thing.__qualname__
+
+        known_thing_names = [t.__name__ for t in known_things]
+
+        new_markdown = []
+        for line_index, line in enumerate(markdown.splitlines(keepends=True)):
+            # Can't convert `this` to `[this][]` in headers, otherwise they break.
+            if line.lstrip().startswith("#"):
+                new_markdown.append(line)
+                continue
+
+            matches = re.findall(r"`([^`]+)`", line)
+
+            for match in matches:
+                thing_name = match
+                if "." not in thing_name and thing_name in known_thing_names:
+                    thing = known_things[known_thing_names.index(thing_name)]
+                else:
+                    thing = _try_import_thing(thing_name)
+
+                if thing is None:
+                    logger.debug(f"Unable to import {thing_name}, leaving it as-is.")
+                    continue
+
+                new_ref = f"[{thing_name}][{_full_path(thing)}]"
+                logger.info(
+                    f"Replacing `{thing_name}` with {new_ref} in {page.file.abs_src_path}:{line_index}"
+                )
+                line = line.replace(f"`{thing_name}`", new_ref)
+
+            new_markdown.append(line)
+
+        return "".join(new_markdown)
+
+
+@functools.cache
+def _try_import_thing(thing: str):
+    try:
+        return import_object(thing)
+    except Exception:
+        return None

project/utils/autoref_plugin_test.py

@@ -0,0 +1,55 @@
+import pytest
+from mkdocs.config.defaults import MkDocsConfig
+from mkdocs.structure.files import File, Files
+from mkdocs.structure.pages import Page
+
+from .autoref_plugin import CustomAutoRefPlugin
+
+
+@pytest.mark.parametrize(
+    ("input", "expected"),
+    [
+        (_header := "## Some header with a ref `lightning.Trainer`", _header),
+        (
+            "a backtick ref: `lightning.Trainer`",
+            "a backtick ref: [lightning.Trainer][lightning.pytorch.trainer.trainer.Trainer]",
+        ),
+        ("`torch.Tensor`", "[torch.Tensor][torch.Tensor]"),
+        (
+            "a proper full ref: "
+            + (
+                _lightning_trainer_ref
+                := "[lightning.Trainer][lightning.pytorch.trainer.trainer.Trainer]"
+            ),
+            # Keep the ref as-is.
+            f"a proper full ref: {_lightning_trainer_ref}",
+        ),
+        ("`foo.bar`", "`foo.bar`"),
+        (
+            "`jax.Array`",
+            # not sure if this will make a proper link in mkdocs though.
+            "[jax.Array][jax.Array]",
+        ),
+        ("`Trainer`", "[Trainer][lightning.pytorch.trainer.trainer.Trainer]"),
+        # since `Trainer` is in the `known_things` list, we add the proper ref.
+    ],
+)
+def test_autoref_plugin(input: str, expected: str):
+    config = MkDocsConfig("mkdocs.yaml")
+    plugin = CustomAutoRefPlugin()
+    result = plugin.on_page_markdown(
+        input,
+        page=Page(
+            title="Test",
+            file=File(
+                "test.md",
+                src_dir="bob",
+                dest_dir="bobo",
+                use_directory_urls=False,
+            ),
+            config=config,
+        ),
+        config=config,
+        files=Files([]),
+    )
+    assert result == expected

pyproject.toml

@@ -35,6 +35,10 @@ dependencies = [
 readme = "README.md"
 requires-python = ">= 3.10"
 
+[project.entry-points."mkdocs.plugins"]
+custom_autoref_plugin = "project.utils.autoref_plugin:CustomAutoRefPlugin"
+
+
 [project.optional-dependencies]
 docs = [
     "mkdocstrings[python]>=0.25.2",