Fix the source link when using a custom basename to include the .py extension

Author: asmeurer

lib/matplotlib/sphinxext/plot_directive.py

@@ -910,7 +910,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,
@@ -970,7 +970,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/test_sphinxext.py

@@ -181,6 +181,28 @@ 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::
+    :image-basename: 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):
     html_dir = tmp_path / '_build' / 'html'
     img_dir = html_dir / '_images'