imporve download role

Author: chrisjsewell

jupyter_sphinx/__init__.py

@@ -19,7 +19,7 @@
     JupyterWidgetViewNode,
     JupyterWidgetStateNode,
     WIDGET_VIEW_MIMETYPE,
-    jupyter_download_role,
+    JupyterDownloadRole,
     CellOutputsToNodes,
 )
 from .execute import JupyterKernel, ExecuteJupyterCells
@@ -272,8 +272,9 @@ def setup(app):
     app.add_directive("jupyter-execute", JupyterCell)
     app.add_directive("jupyter-kernel", JupyterKernel)
     app.add_directive("thebe-button", ThebeButton)
-    app.add_role("jupyter-download:notebook", jupyter_download_role)
-    app.add_role("jupyter-download:script", jupyter_download_role)
+    app.add_role("jupyter-download:notebook", JupyterDownloadRole())
+    app.add_role("jupyter-download:nb", JupyterDownloadRole())
+    app.add_role("jupyter-download:script", JupyterDownloadRole())
     app.add_transform(ExecuteJupyterCells)
     app.add_post_transform(CellOutputsToNodes)
 

jupyter_sphinx/ast.py

@@ -5,8 +5,9 @@
 
 import docutils
 from docutils.parsers.rst import Directive, directives
-from docutils.nodes import math_block, image
+from docutils.nodes import math_block, image, literal
 from sphinx.util import parselinenos
+from sphinx.util.docutils import ReferenceRole
 from sphinx.addnodes import download_reference
 from sphinx.transforms import SphinxTransform
 from sphinx.environment.collectors.asset import ImageCollector
@@ -393,16 +394,20 @@ def attach_outputs(output_nodes, node, thebe_config, cm_language):
         node.children = node.children[::-1]
 
 
-def jupyter_download_role(name, rawtext, text, lineno, inliner):
-    _, filetype = name.split(":")
-    assert filetype in ("notebook", "script")
-    ext = ".ipynb" if filetype == "notebook" else ".py"
-    output_dir = sphinx_abs_dir(inliner.document.settings.env)
-    download_file = text + ext
-    node = download_reference(
-        download_file, download_file, reftarget=os.path.join(output_dir, download_file)
-    )
-    return [node], []
+class JupyterDownloadRole(ReferenceRole):
+    def run(self):
+        _, filetype = self.name.split(":")
+
+        assert filetype in ("notebook", "nb", "script")
+        ext = ".ipynb" if filetype in ("notebook", "nb") else ".py"
+        output_dir = sphinx_abs_dir(self.env)
+        download_file = self.target + ext
+        reftarget=os.path.realpath(os.path.join(output_dir, download_file))
+        node = download_reference(self.rawtext, reftarget=reftarget)
+        self.set_source_info(node)
+        title = self.title if self.has_explicit_title else download_file
+        node += literal(self.rawtext, title, classes=["xref", "download"])
+        return [node], []
 
 
 def get_widgets(notebook):

tests/test_execute.py

@@ -3,10 +3,12 @@
 import os
 import sys
 from io import StringIO
+from unittest.mock import Mock
 
-from sphinx.testing.util import SphinxTestApp, path
+from sphinx.addnodes import download_reference
+from sphinx.testing.util import assert_node, SphinxTestApp, path
 from sphinx.errors import ExtensionError
-from docutils.nodes import raw
+from docutils.nodes import literal, raw
 from nbformat import from_dict
 
 import pytest
@@ -18,6 +20,7 @@
     JupyterWidgetViewNode,
     JupyterWidgetStateNode,
     cell_output_to_nodes,
+    JupyterDownloadRole,
 )
 from jupyter_sphinx.thebelab import ThebeSourceNode, ThebeOutputNode, ThebeButtonNode
 
@@ -564,3 +567,24 @@ def test_image_mimetype_uri(doctree):
         cell = from_dict(cell)
         output_node = cell_output_to_nodes(cell["outputs"], priority, True, output_dir, None)
         assert output_node[0].attributes['uri'] == img_locs[index]
+
+
+@pytest.mark.parametrize('text,reftarget,caption', (
+    ('nb_name', '/jupyter_execute/path/to/nb_name.ipynb', 'nb_name.ipynb'),
+    ('../nb_name', '/jupyter_execute/path/nb_name.ipynb', '../nb_name.ipynb'),
+    ('text <nb_name>', '/jupyter_execute/path/to/nb_name.ipynb', 'text'),
+))
+def test_download_role(text, reftarget, caption):
+    role = JupyterDownloadRole()
+    mock_inliner = Mock()
+    config = {
+        'document.settings.env.app.outdir': 'outdir',
+        'document.settings.env.docname': 'path/to/docname',
+        'document.settings.env.app.srcdir': 'srcdir',
+        'reporter.get_source_and_line': lambda l: ('source', l)
+    }
+    mock_inliner.configure_mock(**config)
+    ret, msg = role('jupyter-download:notebook', text, text, 0, mock_inliner)
+    assert_node(ret[0], [download_reference], reftarget=reftarget)
+    assert_node(ret[0][0], [literal, caption])
+    assert msg == []