Fix toctree sorting

Author: mdboom

cuda_python/docs/exts/release_toc.py

@@ -4,17 +4,24 @@
 from pathlib import Path
 
 from packaging.version import Version
+from sphinx.directives.other import TocTree
 
 
-def sort_release_toctree(app, doctree, docname):
-    """Sort the entries in a release toctree in by version."""
-    if docname == "release":
-        for node in doctree.traverse():
-            if node.tagname == "toctree":
-                node["entries"] = [(Version(Path(x[1]).name.removesuffix("-notes")), x[1]) for x in node["entries"]]
-                node["entries"].sort(key=lambda x: x[0], reverse=True)
-                break
+class TocTreeSorted(TocTree):
+    """A toctree directive that sorts entries by version."""
+
+    def parse_content(self, toctree):
+        super().parse_content(toctree)
+
+        if not toctree["glob"]:
+            return
+
+        toctree["entries"] = [
+            (Version(Path(x[1]).name.removesuffix("-notes")), x[1]) for x in toctree.get("entries", [])
+        ]
+        toctree["entries"].sort(key=lambda x: x[0], reverse=True)
+        toctree["entries"] = [(str(x[0]), x[1]) for x in toctree["entries"]]
 
 
 def setup(app):
-    app.connect("doctree-resolved", sort_release_toctree)
+    app.add_directive("toctree", TocTreeSorted, override=True)