Merge pull request matplotlib#28187 from asmeurer/output-base-name

Add a filename-prefix option to the Sphinx plot directive

Author: timhoffm

lib/matplotlib/sphinxext/plot_directive.py

@@ -47,6 +47,12 @@
 
 The ``.. plot::`` directive supports the following options:
 
+``:filename-prefix:`` : str
+    The base name (without the extension) of the outputted image and script
+    files. The default is to use the same name as the input script, or the
+    name of the RST document if no script is provided. The filename-prefix for
+    each plot directive must be unique.
+
 ``:format:`` : {'python', 'doctest'}
     The format of the input.  If unset, the format is auto-detected.
 
@@ -163,8 +169,10 @@
 be customized by changing the *plot_template*.  See the source of
 :doc:`/api/sphinxext_plot_directive_api` for the templates defined in *TEMPLATE*
 and *TEMPLATE_SRCSET*.
+
 """
 
+from collections import defaultdict
 import contextlib
 import doctest
 from io import StringIO
@@ -182,6 +190,7 @@
 from docutils.parsers.rst.directives.images import Image
 import jinja2  # Sphinx dependency.
 
+from sphinx.environment.collectors import EnvironmentCollector
 from sphinx.errors import ExtensionError
 
 import matplotlib
@@ -265,6 +274,7 @@ class PlotDirective(Directive):
         'scale': directives.nonnegative_int,
         'align': Image.align,
         'class': directives.class_option,
+        'filename-prefix': directives.unchanged,
         'include-source': _option_boolean,
         'show-source-link': _option_boolean,
         'format': _option_format,
@@ -312,9 +322,35 @@ def setup(app):
     app.connect('build-finished', _copy_css_file)
     metadata = {'parallel_read_safe': True, 'parallel_write_safe': True,
                 'version': matplotlib.__version__}
+    app.connect('builder-inited', init_filename_registry)
+    app.add_env_collector(_FilenameCollector)
     return metadata
 
 
+# -----------------------------------------------------------------------------
+# Handle Duplicate Filenames
+# -----------------------------------------------------------------------------
+
+def init_filename_registry(app):
+    env = app.builder.env
+    if not hasattr(env, 'mpl_plot_image_basenames'):
+        env.mpl_plot_image_basenames = defaultdict(set)
+
+
+class _FilenameCollector(EnvironmentCollector):
+    def process_doc(self, app, doctree):
+        pass
+
+    def clear_doc(self, app, env, docname):
+        if docname in env.mpl_plot_image_basenames:
+            del env.mpl_plot_image_basenames[docname]
+
+    def merge_other(self, app, env, docnames, other):
+        for docname in other.mpl_plot_image_basenames:
+            env.mpl_plot_image_basenames[docname].update(
+                other.mpl_plot_image_basenames[docname])
+
+
 # -----------------------------------------------------------------------------
 # Doctest handling
 # -----------------------------------------------------------------------------
@@ -600,6 +636,25 @@ def _parse_srcset(entries):
     return srcset
 
 
+def check_output_base_name(env, output_base):
+    docname = env.docname
+
+    if '.' in output_base or '/' in output_base or '\\' in output_base:
+        raise PlotError(
+            f"The filename-prefix '{output_base}' is invalid. "
+            f"It must not contain dots or slashes.")
+
+    for d in env.mpl_plot_image_basenames:
+        if output_base in env.mpl_plot_image_basenames[d]:
+            if d == docname:
+                raise PlotError(
+                    f"The filename-prefix {output_base!r} is used multiple times.")
+            raise PlotError(f"The filename-prefix {output_base!r} is used multiple"
+                            f"times (it is also used in {env.doc2path(d)}).")
+
+    env.mpl_plot_image_basenames[docname].add(output_base)
+
+
 def render_figures(code, code_path, output_dir, output_base, context,
                    function_name, config, context_reset=False,
                    close_figs=False,
@@ -722,7 +777,8 @@ def render_figures(code, code_path, output_dir, output_base, context,
 
 def run(arguments, content, options, state_machine, state, lineno):
     document = state_machine.document
-    config = document.settings.env.config
+    env = document.settings.env
+    config = env.config
     nofigs = 'nofigs' in options
 
     if config.plot_srcset and setup.app.builder.name == 'singlehtml':
@@ -734,6 +790,7 @@ def run(arguments, content, options, state_machine, state, lineno):
 
     options.setdefault('include-source', config.plot_include_source)
     options.setdefault('show-source-link', config.plot_html_show_source_link)
+    options.setdefault('filename-prefix', None)
 
     if 'class' in options:
         # classes are parsed into a list of string, and output by simply
@@ -775,14 +832,22 @@ def run(arguments, content, options, state_machine, state, lineno):
             function_name = None
 
         code = Path(source_file_name).read_text(encoding='utf-8')
-        output_base = os.path.basename(source_file_name)
+        if options['filename-prefix']:
+            output_base = options['filename-prefix']
+            check_output_base_name(env, output_base)
+        else:
+            output_base = os.path.basename(source_file_name)
     else:
         source_file_name = rst_file
         code = textwrap.dedent("\n".join(map(str, content)))
-        counter = document.attributes.get('_plot_counter', 0) + 1
-        document.attributes['_plot_counter'] = counter
-        base, ext = os.path.splitext(os.path.basename(source_file_name))
-        output_base = '%s-%d.py' % (base, counter)
+        if options['filename-prefix']:
+            output_base = options['filename-prefix']
+            check_output_base_name(env, output_base)
+        else:
+            base, ext = os.path.splitext(os.path.basename(source_file_name))
+            counter = document.attributes.get('_plot_counter', 0) + 1
+            document.attributes['_plot_counter'] = counter
+            output_base = '%s-%d.py' % (base, counter)
         function_name = None
         caption = options.get('caption', '')
 
@@ -846,7 +911,7 @@ def run(arguments, content, options, state_machine, state, lineno):
 
     # save script (if necessary)
     if options['show-source-link']:
-        Path(build_dir, output_base + source_ext).write_text(
+        Path(build_dir, output_base + (source_ext or '.py')).write_text(
             doctest.script_from_examples(code)
             if source_file_name == rst_file and is_doctest
             else code,
@@ -906,7 +971,7 @@ def run(arguments, content, options, state_machine, state, lineno):
         # Not-None src_name signals the need for a source download in the
         # generated html
         if j == 0 and options['show-source-link']:
-            src_name = output_base + source_ext
+            src_name = output_base + (source_ext or '.py')
         else:
             src_name = None
         if config.plot_srcset:

lib/matplotlib/tests/data/tinypages/.gitignore

@@ -1 +1,3 @@
 _build/
+doctrees/
+plot_directive/

lib/matplotlib/tests/data/tinypages/some_plots.rst

@@ -179,3 +179,22 @@ Plot 21 is generated via an include directive:
 Plot 22 uses a different specific function in a file with plot commands:
 
 .. plot:: range6.py range10
+
+Plots 23--25 use filename-prefix.
+
+.. plot::
+   :filename-prefix: custom-basename-6
+
+   plt.plot(range(6))
+
+.. plot:: range4.py
+   :filename-prefix: custom-basename-4
+
+.. plot::
+   :filename-prefix: custom-basename-4-6
+
+   plt.figure()
+   plt.plot(range(4))
+
+   plt.figure()
+   plt.plot(range(6))

lib/matplotlib/tests/test_sphinxext.py

@@ -40,7 +40,9 @@ def build_sphinx_html(source_dir, doctree_dir, html_dir, extra_args=None):
 
 
 def test_tinypages(tmp_path):
-    shutil.copytree(tinypages, tmp_path, dirs_exist_ok=True)
+    shutil.copytree(tinypages, tmp_path, dirs_exist_ok=True,
+                    ignore=shutil.ignore_patterns('_build', 'doctrees',
+                                                   'plot_directive'))
     html_dir = tmp_path / '_build' / 'html'
     img_dir = html_dir / '_images'
     doctree_dir = tmp_path / 'doctrees'
@@ -92,6 +94,11 @@ def plot_directive_file(num):
     assert filecmp.cmp(range_6, plot_file(17))
     # plot 22 is from the range6.py file again, but a different function
     assert filecmp.cmp(range_10, img_dir / 'range6_range10.png')
+    # plots 23--25 use a custom basename
+    assert filecmp.cmp(range_6, img_dir / 'custom-basename-6.png')
+    assert filecmp.cmp(range_4, img_dir / 'custom-basename-4.png')
+    assert filecmp.cmp(range_4, img_dir / 'custom-basename-4-6_00.png')
+    assert filecmp.cmp(range_6, img_dir / 'custom-basename-4-6_01.png')
 
     # Modify the included plot
     contents = (tmp_path / 'included_plot_21.rst').read_bytes()
@@ -176,12 +183,37 @@ def test_show_source_link_false(tmp_path, plot_html_show_source_link):
     assert len(list(html_dir.glob("**/index-1.py"))) == 0
 
 
+def test_plot_html_show_source_link_custom_basename(tmp_path):
+    # Test that source link filename includes .py extension when using custom basename
+    shutil.copyfile(tinypages / 'conf.py', tmp_path / 'conf.py')
+    shutil.copytree(tinypages / '_static', tmp_path / '_static')
+    doctree_dir = tmp_path / 'doctrees'
+    (tmp_path / 'index.rst').write_text("""
+.. plot::
+    :filename-prefix: custom-name
+
+    plt.plot(range(2))
+""")
+    html_dir = tmp_path / '_build' / 'html'
+    build_sphinx_html(tmp_path, doctree_dir, html_dir)
+
+    # Check that source file with .py extension is generated
+    assert len(list(html_dir.glob("**/custom-name.py"))) == 1
+
+    # Check that the HTML contains the correct link with .py extension
+    html_content = (html_dir / 'index.html').read_text()
+    assert 'custom-name.py' in html_content
+
+
 def test_srcset_version(tmp_path):
+    shutil.copytree(tinypages, tmp_path, dirs_exist_ok=True,
+                    ignore=shutil.ignore_patterns('_build', 'doctrees',
+                                                   'plot_directive'))
     html_dir = tmp_path / '_build' / 'html'
     img_dir = html_dir / '_images'
     doctree_dir = tmp_path / 'doctrees'
 
-    build_sphinx_html(tinypages, doctree_dir, html_dir,
+    build_sphinx_html(tmp_path, doctree_dir, html_dir,
                       extra_args=['-D', 'plot_srcset=2x'])
 
     def plot_file(num, suff=''):