diff --git a/docs/macros.py b/docs/macros.py index 23368429..5ca0a148 100644 --- a/docs/macros.py +++ b/docs/macros.py @@ -8,9 +8,27 @@ from pathlib import Path from typing import Any +import torch + if typing.TYPE_CHECKING: from mkdocs_macros.plugin import MacrosPlugin +import lightning + +try: + from mkdocs_autoref_plugin.autoref_plugin import default_reference_sources + + default_reference_sources.extend( + [ + lightning.Trainer, + lightning.LightningModule, + lightning.LightningDataModule, + torch.nn.Module, + ] + ) +except ImportError: + pass + logger = logging.getLogger(__name__) diff --git a/project/datamodules/__init__.py b/project/datamodules/__init__.py index 65bb8580..bc1b12d3 100644 --- a/project/datamodules/__init__.py +++ b/project/datamodules/__init__.py @@ -1,6 +1,6 @@ """Datamodules (datasets + preprocessing + dataloading) -See the :ref:`lightning.LightningDataModule` class for more information. +See the `lightning.LightningDataModule` class for more information. """ from .image_classification import ImageClassificationDataModule diff --git a/project/datamodules/image_classification/imagenet.py b/project/datamodules/image_classification/imagenet.py index 9c774262..769dc9c3 100644 --- a/project/datamodules/image_classification/imagenet.py +++ b/project/datamodules/image_classification/imagenet.py @@ -46,10 +46,10 @@ class ImageNetDataModule(ImageClassificationDataModule): - Made this a subclass of VisionDataModule Notes: + - train_dataloader uses the train split of imagenet2012 and puts away a portion of it for the validation split. - val_dataloader uses the part of the train split of imagenet2012 that was not used for training via `num_imgs_per_val_class` - - TODO: needs to pass split='val' to UnlabeledImagenet. - test_dataloader uses the validation split of imagenet2012 for testing. - TODO: need to pass num_imgs_per_class=-1 for test dataset and split="test". """ diff --git a/project/utils/autoref_plugin.py b/project/utils/autoref_plugin.py deleted file mode 100644 index b706c9bb..00000000 --- a/project/utils/autoref_plugin.py +++ /dev/null @@ -1,216 +0,0 @@ -"""A plugin for the mkdocs documentation engine to provide better support for IDE-friendly links. - -IDEA: Tweak the AutoRefsPlugin so that text in backticks like `this` (more IDE-friendly) are -considered refs when possible. - -TODO: Move to a separate package? -""" - -import functools -import importlib -import inspect -import re -import types - -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 - -# Same as in the mkdocs_autorefs plugin. -logger = get_plugin_logger(__name__) - -default_reference_sources = [ - lightning.Trainer, - lightning.LightningModule, - lightning.LightningDataModule, - torch.nn.Module, -] -"""These are some "known objects" that can be referenced with backticks anywhere in the docs. - -Additionally, if there were modules in here, then any of their public members can also be -referenced. -""" -from mkdocstrings.plugin import MkdocstringsPlugin # noqa -from mkdocstrings_handlers.python.handler import PythonHandler # noqa - - -class CustomAutoRefPlugin(BasePlugin): - """Small mkdocs plugin that converts backticks to refs when possible.""" - - def __init__(self): - super().__init__() - self.default_reference_sources = sum(map(_expand, default_reference_sources), []) - - 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][] - - # TODO: The idea here is to also make all the members of a module referentiable with - # backticks in the same module. The problem is that the "reference" page we create with - # mkdocstrings only contains a simple `::: project.path.to.module` and doesn't have any - # text, so we can't just replace the `backticks` with refs, since mkdocstrings hasn't yet - # processed the module into a page with the reference docs. This seems to be happening - # in a markdown extension (the `MkdocstringsExtension`). - - # file = page.file.abs_src_path - # if file and "reference/project" in file: - # relative_path = file[file.index("reference/") :].removeprefix("reference/") - # module_path = relative_path.replace("/", ".").replace(".md", "") - # if module_path.endswith(".index"): - # module_path = module_path.removesuffix(".index") - # logger.error( - # f"file {relative_path} is the reference page for the python module {module_path}" - # ) - # if "algorithms/example" in file: - # assert False, markdown - # additional_objects = _expand(module_path) - if referenced_packages := page.meta.get("additional_python_references", []): - logger.debug(f"Loading extra references: {referenced_packages}") - additional_objects: list[object] = _get_referencable_objects_from_doc_page_header( - referenced_packages - ) - else: - additional_objects = [] - - if additional_objects: - additional_objects = [ - obj - for obj in additional_objects - if ( - inspect.isfunction(obj) - or inspect.isclass(obj) - or inspect.ismodule(obj) - or inspect.ismethod(obj) - ) - # and (hasattr(obj, "__name__") or hasattr(obj, "__qualname__")) - ] - - known_objects_for_this_module = self.default_reference_sources + additional_objects - known_object_names = [t.__name__ for t in known_objects_for_this_module] - - new_markdown = [] - # TODO: This changes things inside code blocks, which is not desired! - in_code_block = False - - 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 - if "```" in line: - in_code_block = not in_code_block - if in_code_block: - new_markdown.append(line) - continue - - matches = re.findall(r"`([^`]+)`", line) - for match in matches: - thing_name = match - if any(char in thing_name for char in ["/", " ", "-"]): - continue - if thing_name in known_object_names: - # References like `JaxTrainer` (which are in a module that we're aware of). - thing = known_objects_for_this_module[known_object_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.debug( - 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) - - -def _expand(obj: types.ModuleType | object) -> list[object]: - if not inspect.ismodule(obj): - # The ref is something else (a class, function, etc.) - return [obj] - - # The ref is a package, so we import everything from it. - # equivalent of `from package import *` - if hasattr(obj, "__all__"): - return [getattr(obj, name) for name in obj.__all__] - else: - objects_in_global_scope = [v for k, v in vars(obj).items() if not k.startswith("_")] - # Don't consider any external modules that were imported in the global scope. - source_file = inspect.getsourcefile(obj) - # too obtuse, but whatever - return [ - v - for v in objects_in_global_scope - if not ( - (inspect.ismodule(v) and getattr(v, "__file__", None) is None) # built-in module. - or (inspect.ismodule(v) and inspect.getsourcefile(v) != source_file) - ) - ] - - -def import_object(target_path: str): - """Imports the object at the given path.""" - - # todo: what is the difference between this here and `hydra.utils.get_object` ? - assert not target_path.endswith( - ".py" - ), "expect a valid python path like 'module.submodule.object'" - if "." not in target_path: - return importlib.import_module(target_path) - - parts = target_path.split(".") - try: - return importlib.import_module(name=f".{parts[-1]}", package=".".join(parts[:-1])) - except (ModuleNotFoundError, AttributeError): - pass - exc = None - for i in range(1, len(parts)): - module_name = ".".join(parts[:i]) - obj_path = parts[i:] - try: - module = importlib.import_module(module_name) - obj = getattr(module, obj_path[0]) - for part in obj_path[1:]: - obj = getattr(obj, part) - return obj - except (ModuleNotFoundError, AttributeError) as _exc: - exc = _exc - continue - assert exc is not None - raise ModuleNotFoundError(f"Unable to import the {target_path=}!") from exc - - -def _get_referencable_objects_from_doc_page_header(doc_page_references: list[str]): - additional_objects: list[object] = [] - for package in doc_page_references: - additional_ref_source = import_object(package) - additional_objects.extend(_expand(additional_ref_source)) - return additional_objects - - -def _full_path(thing) -> str: - if inspect.ismodule(thing): - return thing.__name__ - return thing.__module__ + "." + getattr(thing, "__qualname__", thing.__name__) - - -@functools.cache -def _try_import_thing(thing: str): - try: - return import_object(thing) - except Exception: - return None diff --git a/project/utils/autoref_plugin_test.py b/project/utils/autoref_plugin_test.py deleted file mode 100644 index feec2688..00000000 --- a/project/utils/autoref_plugin_test.py +++ /dev/null @@ -1,82 +0,0 @@ -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. - ("`.devcontainer/devcontainer.json`", "`.devcontainer/devcontainer.json`"), - ], -) -def test_autoref_plugin(input: str, expected: str): - config: MkDocsConfig = MkDocsConfig("mkdocs.yaml") # type: ignore (weird!) - 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 - - -def test_ref_using_additional_python_references(): - mkdocs_config: MkDocsConfig = MkDocsConfig("mkdocs.yaml") # type: ignore (weird!) - - plugin = CustomAutoRefPlugin() - - page = Page( - title="Test", - file=File( - "test.md", - src_dir="bob", - dest_dir="bobo", - use_directory_urls=False, - ), - config=mkdocs_config, - ) - page.meta = {"additional_python_references": ["project.algorithms.image_classifier"]} - - result = plugin.on_page_markdown( - "`ImageClassifier`", - page=page, - config=mkdocs_config, - files=Files([]), - ) - assert result == "[`ImageClassifier`][project.algorithms.image_classifier.ImageClassifier]" diff --git a/pyproject.toml b/pyproject.toml index 88d68c73..9d5f8c23 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,9 +45,6 @@ dependencies = [ readme = "README.md" requires-python = ">= 3.10" -[project.entry-points."mkdocs.plugins"] -custom_autoref_plugin = "project.utils.autoref_plugin:CustomAutoRefPlugin" - [dependency-groups] dev = [ "mktestdocs>=0.2.2", @@ -75,6 +72,7 @@ docs = [ "mkdocs-video>=1.5.0", "mkdocs-section-index>=0.3.9", "mkdocs-macros-plugin>=1.0.5", + "mkdocs-autoref-plugin", ] gpu = ["jax[cuda12]>=0.4.31"] @@ -118,6 +116,7 @@ managed = true [tool.uv.sources] remote-slurm-executor = { git = "https://github.com/lebrice/remote-slurm-executor", branch = "master" } +mkdocs-autoref-plugin = { git = "https://github.com/lebrice/mkdocs-autoref-plugin" } # [tool.uv.tasks] # serve-docs = "mkdocs serve" diff --git a/uv.lock b/uv.lock index b002472a..d4f1266e 100644 --- a/uv.lock +++ b/uv.lock @@ -2097,6 +2097,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/22/5b/dbc6a8cddc9cfa9c4971d59fb12bb8d42e161b7e7f8cc89e49137c5b279c/mkdocs-1.6.1-py3-none-any.whl", hash = "sha256:db91759624d1647f3f34aa0c3f327dd2601beae39a366d6e064c03468d35c20e", size = 3864451 }, ] +[[package]] +name = "mkdocs-autoref-plugin" +version = "0.0.1" +source = { git = "https://github.com/lebrice/mkdocs-autoref-plugin#92e6d2f8d07d2b79317a0222c89714fd06e8330a" } +dependencies = [ + { name = "mkdocs-autorefs" }, + { name = "mkdocstrings", extra = ["python"] }, +] + [[package]] name = "mkdocs-autorefs" version = "1.2.0" @@ -2241,7 +2250,7 @@ wheels = [ [[package]] name = "mkdocstrings" -version = "0.26.2" +version = "0.27.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -2253,9 +2262,9 @@ dependencies = [ { name = "platformdirs" }, { name = "pymdown-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/c0/76/0475d10d27f3384df3a6ddfdf4a4fdfef83766f77cd4e327d905dc956c15/mkdocstrings-0.26.2.tar.gz", hash = "sha256:34a8b50f1e6cfd29546c6c09fbe02154adfb0b361bb758834bf56aa284ba876e", size = 92512 } +sdist = { url = "https://files.pythonhosted.org/packages/e2/5a/5de70538c2cefae7ac3a15b5601e306ef3717290cb2aab11d51cbbc2d1c0/mkdocstrings-0.27.0.tar.gz", hash = "sha256:16adca6d6b0a1f9e0c07ff0b02ced8e16f228a9d65a37c063ec4c14d7b76a657", size = 94830 } wheels = [ - { url = "https://files.pythonhosted.org/packages/80/b6/4ee320d7c313da3774eff225875eb278f7e6bb26a9cd8e680b8dbc38fdea/mkdocstrings-0.26.2-py3-none-any.whl", hash = "sha256:1248f3228464f3b8d1a15bd91249ce1701fe3104ac517a5f167a0e01ca850ba5", size = 29716 }, + { url = "https://files.pythonhosted.org/packages/cd/10/4c27c3063c2b3681a4b7942f8dbdeb4fa34fecb2c19b594e7345ebf4f86f/mkdocstrings-0.27.0-py3-none-any.whl", hash = "sha256:6ceaa7ea830770959b55a16203ac63da24badd71325b96af950e59fd37366332", size = 30658 }, ] [package.optional-dependencies] @@ -3890,6 +3899,7 @@ dependencies = [ [package.optional-dependencies] docs = [ { name = "black" }, + { name = "mkdocs-autoref-plugin" }, { name = "mkdocs-awesome-pages-plugin" }, { name = "mkdocs-gen-files" }, { name = "mkdocs-literate-nav" }, @@ -3940,6 +3950,7 @@ requires-dist = [ { name = "lightning", specifier = ">=2.4.0" }, { name = "matplotlib", specifier = ">=3.9.2" }, { name = "mkdocs", specifier = ">=1.6.1" }, + { name = "mkdocs-autoref-plugin", marker = "extra == 'docs'", git = "https://github.com/lebrice/mkdocs-autoref-plugin" }, { name = "mkdocs-awesome-pages-plugin", marker = "extra == 'docs'", specifier = ">=2.9.3" }, { name = "mkdocs-gen-files", marker = "extra == 'docs'", specifier = ">=0.5.0" }, { name = "mkdocs-literate-nav", marker = "extra == 'docs'", specifier = ">=0.6.1" },