From a8017583d0dfad360a56e44da8396abfc1291816 Mon Sep 17 00:00:00 2001
From: Bruno Beltran <brunobeltran0@gmail.com>
Date: Fri, 4 Dec 2020 10:34:47 -0800
Subject: [PATCH] BF: plot_directive template to allow class+caption

---
 lib/matplotlib/sphinxext/plot_directive.py    | 15 +++++++-------
 lib/matplotlib/tests/test_sphinxext.py        |  7 ++++++-
 lib/matplotlib/tests/tinypages/some_plots.rst | 20 +++++++++++++++++++
 3 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/lib/matplotlib/sphinxext/plot_directive.py b/lib/matplotlib/sphinxext/plot_directive.py
index 7de243f2ead6..1b031ff01a0b 100644
--- a/lib/matplotlib/sphinxext/plot_directive.py
+++ b/lib/matplotlib/sphinxext/plot_directive.py
@@ -344,7 +344,6 @@ def split_code_at_show(text):
 # Template
 # -----------------------------------------------------------------------------
 
-
 TEMPLATE = """
 {{ source_code }}
 
@@ -379,9 +378,8 @@ def split_code_at_show(text):
         `{{ fmt }} <{{ dest_dir }}/{{ img.basename }}.{{ fmt }}>`__
         {%- endfor -%}
         )
-      {%- endif -%}
-
-      {{ caption }}
+      {%- endif %}
+{{ caption }}
    {% endfor %}
 
 .. only:: not html
@@ -392,7 +390,7 @@ def split_code_at_show(text):
       {{ option }}
       {% endfor %}
 
-      {{ caption }}
+{{ caption }}
    {% endfor %}
 
 """
@@ -632,9 +630,12 @@ def run(arguments, content, options, state_machine, state, lineno):
 
     options.setdefault('include-source', config.plot_include_source)
     if 'class' in options:
-        options['class'] = options['class'] + ' plot-directive'
+        # classes are parsed into a list of string, and output by simply
+        # printing the list, abusing the fact that RST guarantees to strip
+        # non-conforming characters
+        options['class'] = ['plot-directive'] + options['class']
     else:
-        options.setdefault('class', 'plot-directive')
+        options.setdefault('class', ['plot-directive'])
     keep_context = 'context' in options
     context_opt = None if not keep_context else options['context']
 
diff --git a/lib/matplotlib/tests/test_sphinxext.py b/lib/matplotlib/tests/test_sphinxext.py
index 3ce8d2951d07..cc03a87270f9 100644
--- a/lib/matplotlib/tests/test_sphinxext.py
+++ b/lib/matplotlib/tests/test_sphinxext.py
@@ -11,7 +11,8 @@
 pytest.importorskip('sphinx')
 
 
-def test_tinypages(tmpdir):
+def test_tinypages():
+    tmpdir = '/tmp/bb'
     tmp_path = Path(tmpdir)
     html_dir = tmp_path / 'html'
     doctree_dir = tmp_path / 'doctrees'
@@ -57,3 +58,7 @@ def plot_file(num):
     assert b'Plot 17 uses the caption option.' in html_contents
     # check if figure caption made it into html file
     assert b'This is the caption for plot 18.' in html_contents
+    # check if the custom classes made it into the html file
+    assert b'plot-directive my-class my-other-class' in html_contents
+    # check that the multi-image caption is applied twice
+    assert html_contents.count(b'This caption applies to both plots.') == 2
diff --git a/lib/matplotlib/tests/tinypages/some_plots.rst b/lib/matplotlib/tests/tinypages/some_plots.rst
index 5a71abc9e761..514552decfee 100644
--- a/lib/matplotlib/tests/tinypages/some_plots.rst
+++ b/lib/matplotlib/tests/tinypages/some_plots.rst
@@ -141,3 +141,23 @@ using the :caption: option:
 
 .. plot:: range4.py
    :caption: This is the caption for plot 18.
+
+Plot 19 uses shows that the "plot-directive" class is still appended, even if
+we request other custom classes:
+
+.. plot:: range4.py
+   :class: my-class my-other-class
+
+    Should also have a caption.
+
+Plot 20 shows that the default template correctly prints the multi-image
+scenario:
+
+.. plot::
+   :caption: This caption applies to both plots.
+
+   plt.figure()
+   plt.plot(range(6))
+
+   plt.figure()
+   plt.plot(range(4))