plot_directive.py 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933
  1. """
  2. A directive for including a Matplotlib plot in a Sphinx document
  3. ================================================================
  4. This is a Sphinx extension providing a reStructuredText directive
  5. ``.. plot::`` for including a plot in a Sphinx document.
  6. In HTML output, ``.. plot::`` will include a .png file with a link
  7. to a high-res .png and .pdf. In LaTeX output, it will include a .pdf.
  8. The plot content may be defined in one of three ways:
  9. 1. **A path to a source file** as the argument to the directive::
  10. .. plot:: path/to/plot.py
  11. When a path to a source file is given, the content of the
  12. directive may optionally contain a caption for the plot::
  13. .. plot:: path/to/plot.py
  14. The plot caption.
  15. Additionally, one may specify the name of a function to call (with
  16. no arguments) immediately after importing the module::
  17. .. plot:: path/to/plot.py plot_function1
  18. 2. Included as **inline content** to the directive::
  19. .. plot::
  20. import matplotlib.pyplot as plt
  21. plt.plot([1, 2, 3], [4, 5, 6])
  22. plt.title("A plotting exammple")
  23. 3. Using **doctest** syntax::
  24. .. plot::
  25. A plotting example:
  26. >>> import matplotlib.pyplot as plt
  27. >>> plt.plot([1, 2, 3], [4, 5, 6])
  28. Options
  29. -------
  30. The ``.. plot::`` directive supports the following options:
  31. ``:format:`` : {'python', 'doctest'}
  32. The format of the input. If unset, the format is auto-detected.
  33. ``:include-source:`` : bool
  34. Whether to display the source code. The default can be changed using
  35. the ``plot_include_source`` variable in :file:`conf.py` (which itself
  36. defaults to False).
  37. ``:show-source-link:`` : bool
  38. Whether to show a link to the source in HTML. The default can be
  39. changed using the ``plot_html_show_source_link`` variable in
  40. :file:`conf.py` (which itself defaults to True).
  41. ``:context:`` : bool or str
  42. If provided, the code will be run in the context of all previous plot
  43. directives for which the ``:context:`` option was specified. This only
  44. applies to inline code plot directives, not those run from files. If
  45. the ``:context: reset`` option is specified, the context is reset
  46. for this and future plots, and previous figures are closed prior to
  47. running the code. ``:context: close-figs`` keeps the context but closes
  48. previous figures before running the code.
  49. ``:nofigs:`` : bool
  50. If specified, the code block will be run, but no figures will be
  51. inserted. This is usually useful with the ``:context:`` option.
  52. ``:caption:`` : str
  53. If specified, the option's argument will be used as a caption for the
  54. figure. This overwrites the caption given in the content, when the plot
  55. is generated from a file.
  56. Additionally, this directive supports all the options of the `image directive
  57. <https://docutils.sourceforge.io/docs/ref/rst/directives.html#image>`_,
  58. except for ``:target:`` (since plot will add its own target). These include
  59. ``:alt:``, ``:height:``, ``:width:``, ``:scale:``, ``:align:`` and ``:class:``.
  60. Configuration options
  61. ---------------------
  62. The plot directive has the following configuration options:
  63. plot_include_source
  64. Default value for the include-source option (default: False).
  65. plot_html_show_source_link
  66. Whether to show a link to the source in HTML (default: True).
  67. plot_pre_code
  68. Code that should be executed before each plot. If None (the default),
  69. it will default to a string containing::
  70. import numpy as np
  71. from matplotlib import pyplot as plt
  72. plot_basedir
  73. Base directory, to which ``plot::`` file names are relative to.
  74. If None or empty (the default), file names are relative to the
  75. directory where the file containing the directive is.
  76. plot_formats
  77. File formats to generate (default: ['png', 'hires.png', 'pdf']).
  78. List of tuples or strings::
  79. [(suffix, dpi), suffix, ...]
  80. that determine the file format and the DPI. For entries whose
  81. DPI was omitted, sensible defaults are chosen. When passing from
  82. the command line through sphinx_build the list should be passed as
  83. suffix:dpi,suffix:dpi, ...
  84. plot_html_show_formats
  85. Whether to show links to the files in HTML (default: True).
  86. plot_rcparams
  87. A dictionary containing any non-standard rcParams that should
  88. be applied before each plot (default: {}).
  89. plot_apply_rcparams
  90. By default, rcParams are applied when ``:context:`` option is not used
  91. in a plot directive. If set, this configuration option overrides this
  92. behavior and applies rcParams before each plot.
  93. plot_working_directory
  94. By default, the working directory will be changed to the directory of
  95. the example, so the code can get at its data files, if any. Also its
  96. path will be added to `sys.path` so it can import any helper modules
  97. sitting beside it. This configuration option can be used to specify
  98. a central directory (also added to `sys.path`) where data files and
  99. helper modules for all code are located.
  100. plot_template
  101. Provide a customized template for preparing restructured text.
  102. plot_srcset
  103. Allow the srcset image option for responsive image resolutions. List of
  104. strings with the multiplicative factors followed by an "x".
  105. e.g. ["2.0x", "1.5x"]. "2.0x" will create a png with the default "png"
  106. resolution from plot_formats, multiplied by 2. If plot_srcset is
  107. specified, the plot directive uses the
  108. :doc:`/api/sphinxext_figmpl_directive_api` (instead of the usual figure
  109. directive) in the intermediary rst file that is generated.
  110. The plot_srcset option is incompatible with *singlehtml* builds, and an
  111. error will be raised.
  112. Notes on how it works
  113. ---------------------
  114. The plot directive runs the code it is given, either in the source file or the
  115. code under the directive. The figure created (if any) is saved in the sphinx
  116. build directory under a subdirectory named ``plot_directive``. It then creates
  117. an intermediate rst file that calls a ``.. figure:`` directive (or
  118. ``.. figmpl::`` directive if ``plot_srcset`` is being used) and has links to
  119. the ``*.png`` files in the ``plot_directive`` directory. These translations can
  120. be customized by changing the *plot_template*. See the source of
  121. :doc:`/api/sphinxext_plot_directive_api` for the templates defined in *TEMPLATE*
  122. and *TEMPLATE_SRCSET*.
  123. """
  124. import contextlib
  125. import doctest
  126. from io import StringIO
  127. import itertools
  128. import os
  129. from os.path import relpath
  130. from pathlib import Path
  131. import re
  132. import shutil
  133. import sys
  134. import textwrap
  135. import traceback
  136. from docutils.parsers.rst import directives, Directive
  137. from docutils.parsers.rst.directives.images import Image
  138. import jinja2 # Sphinx dependency.
  139. from sphinx.errors import ExtensionError
  140. import matplotlib
  141. from matplotlib.backend_bases import FigureManagerBase
  142. import matplotlib.pyplot as plt
  143. from matplotlib import _pylab_helpers, cbook
  144. matplotlib.use("agg")
  145. __version__ = 2
  146. # -----------------------------------------------------------------------------
  147. # Registration hook
  148. # -----------------------------------------------------------------------------
  149. def _option_boolean(arg):
  150. if not arg or not arg.strip():
  151. # no argument given, assume used as a flag
  152. return True
  153. elif arg.strip().lower() in ('no', '0', 'false'):
  154. return False
  155. elif arg.strip().lower() in ('yes', '1', 'true'):
  156. return True
  157. else:
  158. raise ValueError(f'{arg!r} unknown boolean')
  159. def _option_context(arg):
  160. if arg in [None, 'reset', 'close-figs']:
  161. return arg
  162. raise ValueError("Argument should be None or 'reset' or 'close-figs'")
  163. def _option_format(arg):
  164. return directives.choice(arg, ('python', 'doctest'))
  165. def mark_plot_labels(app, document):
  166. """
  167. To make plots referenceable, we need to move the reference from the
  168. "htmlonly" (or "latexonly") node to the actual figure node itself.
  169. """
  170. for name, explicit in document.nametypes.items():
  171. if not explicit:
  172. continue
  173. labelid = document.nameids[name]
  174. if labelid is None:
  175. continue
  176. node = document.ids[labelid]
  177. if node.tagname in ('html_only', 'latex_only'):
  178. for n in node:
  179. if n.tagname == 'figure':
  180. sectname = name
  181. for c in n:
  182. if c.tagname == 'caption':
  183. sectname = c.astext()
  184. break
  185. node['ids'].remove(labelid)
  186. node['names'].remove(name)
  187. n['ids'].append(labelid)
  188. n['names'].append(name)
  189. document.settings.env.labels[name] = \
  190. document.settings.env.docname, labelid, sectname
  191. break
  192. class PlotDirective(Directive):
  193. """The ``.. plot::`` directive, as documented in the module's docstring."""
  194. has_content = True
  195. required_arguments = 0
  196. optional_arguments = 2
  197. final_argument_whitespace = False
  198. option_spec = {
  199. 'alt': directives.unchanged,
  200. 'height': directives.length_or_unitless,
  201. 'width': directives.length_or_percentage_or_unitless,
  202. 'scale': directives.nonnegative_int,
  203. 'align': Image.align,
  204. 'class': directives.class_option,
  205. 'include-source': _option_boolean,
  206. 'show-source-link': _option_boolean,
  207. 'format': _option_format,
  208. 'context': _option_context,
  209. 'nofigs': directives.flag,
  210. 'caption': directives.unchanged,
  211. }
  212. def run(self):
  213. """Run the plot directive."""
  214. try:
  215. return run(self.arguments, self.content, self.options,
  216. self.state_machine, self.state, self.lineno)
  217. except Exception as e:
  218. raise self.error(str(e))
  219. def _copy_css_file(app, exc):
  220. if exc is None and app.builder.format == 'html':
  221. src = cbook._get_data_path('plot_directive/plot_directive.css')
  222. dst = app.outdir / Path('_static')
  223. dst.mkdir(exist_ok=True)
  224. # Use copyfile because we do not want to copy src's permissions.
  225. shutil.copyfile(src, dst / Path('plot_directive.css'))
  226. def setup(app):
  227. setup.app = app
  228. setup.config = app.config
  229. setup.confdir = app.confdir
  230. app.add_directive('plot', PlotDirective)
  231. app.add_config_value('plot_pre_code', None, True)
  232. app.add_config_value('plot_include_source', False, True)
  233. app.add_config_value('plot_html_show_source_link', True, True)
  234. app.add_config_value('plot_formats', ['png', 'hires.png', 'pdf'], True)
  235. app.add_config_value('plot_basedir', None, True)
  236. app.add_config_value('plot_html_show_formats', True, True)
  237. app.add_config_value('plot_rcparams', {}, True)
  238. app.add_config_value('plot_apply_rcparams', False, True)
  239. app.add_config_value('plot_working_directory', None, True)
  240. app.add_config_value('plot_template', None, True)
  241. app.add_config_value('plot_srcset', [], True)
  242. app.connect('doctree-read', mark_plot_labels)
  243. app.add_css_file('plot_directive.css')
  244. app.connect('build-finished', _copy_css_file)
  245. metadata = {'parallel_read_safe': True, 'parallel_write_safe': True,
  246. 'version': matplotlib.__version__}
  247. return metadata
  248. # -----------------------------------------------------------------------------
  249. # Doctest handling
  250. # -----------------------------------------------------------------------------
  251. def contains_doctest(text):
  252. try:
  253. # check if it's valid Python as-is
  254. compile(text, '<string>', 'exec')
  255. return False
  256. except SyntaxError:
  257. pass
  258. r = re.compile(r'^\s*>>>', re.M)
  259. m = r.search(text)
  260. return bool(m)
  261. def _split_code_at_show(text, function_name):
  262. """Split code at plt.show()."""
  263. is_doctest = contains_doctest(text)
  264. if function_name is None:
  265. parts = []
  266. part = []
  267. for line in text.split("\n"):
  268. if ((not is_doctest and line.startswith('plt.show(')) or
  269. (is_doctest and line.strip() == '>>> plt.show()')):
  270. part.append(line)
  271. parts.append("\n".join(part))
  272. part = []
  273. else:
  274. part.append(line)
  275. if "\n".join(part).strip():
  276. parts.append("\n".join(part))
  277. else:
  278. parts = [text]
  279. return is_doctest, parts
  280. # -----------------------------------------------------------------------------
  281. # Template
  282. # -----------------------------------------------------------------------------
  283. _SOURCECODE = """
  284. {{ source_code }}
  285. .. only:: html
  286. {% if src_name or (html_show_formats and not multi_image) %}
  287. (
  288. {%- if src_name -%}
  289. :download:`Source code <{{ build_dir }}/{{ src_name }}>`
  290. {%- endif -%}
  291. {%- if html_show_formats and not multi_image -%}
  292. {%- for img in images -%}
  293. {%- for fmt in img.formats -%}
  294. {%- if src_name or not loop.first -%}, {% endif -%}
  295. :download:`{{ fmt }} <{{ build_dir }}/{{ img.basename }}.{{ fmt }}>`
  296. {%- endfor -%}
  297. {%- endfor -%}
  298. {%- endif -%}
  299. )
  300. {% endif %}
  301. """
  302. TEMPLATE_SRCSET = _SOURCECODE + """
  303. {% for img in images %}
  304. .. figure-mpl:: {{ build_dir }}/{{ img.basename }}.{{ default_fmt }}
  305. {% for option in options -%}
  306. {{ option }}
  307. {% endfor %}
  308. {%- if caption -%}
  309. {{ caption }} {# appropriate leading whitespace added beforehand #}
  310. {% endif -%}
  311. {%- if srcset -%}
  312. :srcset: {{ build_dir }}/{{ img.basename }}.{{ default_fmt }}
  313. {%- for sr in srcset -%}
  314. , {{ build_dir }}/{{ img.basename }}.{{ sr }}.{{ default_fmt }} {{sr}}
  315. {%- endfor -%}
  316. {% endif %}
  317. {% if html_show_formats and multi_image %}
  318. (
  319. {%- for fmt in img.formats -%}
  320. {%- if not loop.first -%}, {% endif -%}
  321. :download:`{{ fmt }} <{{ build_dir }}/{{ img.basename }}.{{ fmt }}>`
  322. {%- endfor -%}
  323. )
  324. {% endif %}
  325. {% endfor %}
  326. .. only:: not html
  327. {% for img in images %}
  328. .. figure-mpl:: {{ build_dir }}/{{ img.basename }}.*
  329. {% for option in options -%}
  330. {{ option }}
  331. {% endfor -%}
  332. {{ caption }} {# appropriate leading whitespace added beforehand #}
  333. {% endfor %}
  334. """
  335. TEMPLATE = _SOURCECODE + """
  336. {% for img in images %}
  337. .. figure:: {{ build_dir }}/{{ img.basename }}.{{ default_fmt }}
  338. {% for option in options -%}
  339. {{ option }}
  340. {% endfor %}
  341. {% if html_show_formats and multi_image -%}
  342. (
  343. {%- for fmt in img.formats -%}
  344. {%- if not loop.first -%}, {% endif -%}
  345. :download:`{{ fmt }} <{{ build_dir }}/{{ img.basename }}.{{ fmt }}>`
  346. {%- endfor -%}
  347. )
  348. {%- endif -%}
  349. {{ caption }} {# appropriate leading whitespace added beforehand #}
  350. {% endfor %}
  351. .. only:: not html
  352. {% for img in images %}
  353. .. figure:: {{ build_dir }}/{{ img.basename }}.*
  354. {% for option in options -%}
  355. {{ option }}
  356. {% endfor -%}
  357. {{ caption }} {# appropriate leading whitespace added beforehand #}
  358. {% endfor %}
  359. """
  360. exception_template = """
  361. .. only:: html
  362. [`source code <%(linkdir)s/%(basename)s.py>`__]
  363. Exception occurred rendering plot.
  364. """
  365. # the context of the plot for all directives specified with the
  366. # :context: option
  367. plot_context = dict()
  368. class ImageFile:
  369. def __init__(self, basename, dirname):
  370. self.basename = basename
  371. self.dirname = dirname
  372. self.formats = []
  373. def filename(self, format):
  374. return os.path.join(self.dirname, f"{self.basename}.{format}")
  375. def filenames(self):
  376. return [self.filename(fmt) for fmt in self.formats]
  377. def out_of_date(original, derived, includes=None):
  378. """
  379. Return whether *derived* is out-of-date relative to *original* or any of
  380. the RST files included in it using the RST include directive (*includes*).
  381. *derived* and *original* are full paths, and *includes* is optionally a
  382. list of full paths which may have been included in the *original*.
  383. """
  384. if not os.path.exists(derived):
  385. return True
  386. if includes is None:
  387. includes = []
  388. files_to_check = [original, *includes]
  389. def out_of_date_one(original, derived_mtime):
  390. return (os.path.exists(original) and
  391. derived_mtime < os.stat(original).st_mtime)
  392. derived_mtime = os.stat(derived).st_mtime
  393. return any(out_of_date_one(f, derived_mtime) for f in files_to_check)
  394. class PlotError(RuntimeError):
  395. pass
  396. def _run_code(code, code_path, ns=None, function_name=None):
  397. """
  398. Import a Python module from a path, and run the function given by
  399. name, if function_name is not None.
  400. """
  401. # Change the working directory to the directory of the example, so
  402. # it can get at its data files, if any. Add its path to sys.path
  403. # so it can import any helper modules sitting beside it.
  404. pwd = os.getcwd()
  405. if setup.config.plot_working_directory is not None:
  406. try:
  407. os.chdir(setup.config.plot_working_directory)
  408. except OSError as err:
  409. raise OSError(f'{err}\n`plot_working_directory` option in '
  410. f'Sphinx configuration file must be a valid '
  411. f'directory path') from err
  412. except TypeError as err:
  413. raise TypeError(f'{err}\n`plot_working_directory` option in '
  414. f'Sphinx configuration file must be a string or '
  415. f'None') from err
  416. elif code_path is not None:
  417. dirname = os.path.abspath(os.path.dirname(code_path))
  418. os.chdir(dirname)
  419. with cbook._setattr_cm(
  420. sys, argv=[code_path], path=[os.getcwd(), *sys.path]), \
  421. contextlib.redirect_stdout(StringIO()):
  422. try:
  423. if ns is None:
  424. ns = {}
  425. if not ns:
  426. if setup.config.plot_pre_code is None:
  427. exec('import numpy as np\n'
  428. 'from matplotlib import pyplot as plt\n', ns)
  429. else:
  430. exec(str(setup.config.plot_pre_code), ns)
  431. if "__main__" in code:
  432. ns['__name__'] = '__main__'
  433. # Patch out non-interactive show() to avoid triggering a warning.
  434. with cbook._setattr_cm(FigureManagerBase, show=lambda self: None):
  435. exec(code, ns)
  436. if function_name is not None:
  437. exec(function_name + "()", ns)
  438. except (Exception, SystemExit) as err:
  439. raise PlotError(traceback.format_exc()) from err
  440. finally:
  441. os.chdir(pwd)
  442. return ns
  443. def clear_state(plot_rcparams, close=True):
  444. if close:
  445. plt.close('all')
  446. matplotlib.rc_file_defaults()
  447. matplotlib.rcParams.update(plot_rcparams)
  448. def get_plot_formats(config):
  449. default_dpi = {'png': 80, 'hires.png': 200, 'pdf': 200}
  450. formats = []
  451. plot_formats = config.plot_formats
  452. for fmt in plot_formats:
  453. if isinstance(fmt, str):
  454. if ':' in fmt:
  455. suffix, dpi = fmt.split(':')
  456. formats.append((str(suffix), int(dpi)))
  457. else:
  458. formats.append((fmt, default_dpi.get(fmt, 80)))
  459. elif isinstance(fmt, (tuple, list)) and len(fmt) == 2:
  460. formats.append((str(fmt[0]), int(fmt[1])))
  461. else:
  462. raise PlotError('invalid image format "%r" in plot_formats' % fmt)
  463. return formats
  464. def _parse_srcset(entries):
  465. """
  466. Parse srcset for multiples...
  467. """
  468. srcset = {}
  469. for entry in entries:
  470. entry = entry.strip()
  471. if len(entry) >= 2:
  472. mult = entry[:-1]
  473. srcset[float(mult)] = entry
  474. else:
  475. raise ExtensionError(f'srcset argument {entry!r} is invalid.')
  476. return srcset
  477. def render_figures(code, code_path, output_dir, output_base, context,
  478. function_name, config, context_reset=False,
  479. close_figs=False,
  480. code_includes=None):
  481. """
  482. Run a pyplot script and save the images in *output_dir*.
  483. Save the images under *output_dir* with file names derived from
  484. *output_base*
  485. """
  486. if function_name is not None:
  487. output_base = f'{output_base}_{function_name}'
  488. formats = get_plot_formats(config)
  489. # Try to determine if all images already exist
  490. is_doctest, code_pieces = _split_code_at_show(code, function_name)
  491. # Look for single-figure output files first
  492. img = ImageFile(output_base, output_dir)
  493. for format, dpi in formats:
  494. if context or out_of_date(code_path, img.filename(format),
  495. includes=code_includes):
  496. all_exists = False
  497. break
  498. img.formats.append(format)
  499. else:
  500. all_exists = True
  501. if all_exists:
  502. return [(code, [img])]
  503. # Then look for multi-figure output files
  504. results = []
  505. for i, code_piece in enumerate(code_pieces):
  506. images = []
  507. for j in itertools.count():
  508. if len(code_pieces) > 1:
  509. img = ImageFile('%s_%02d_%02d' % (output_base, i, j),
  510. output_dir)
  511. else:
  512. img = ImageFile('%s_%02d' % (output_base, j), output_dir)
  513. for fmt, dpi in formats:
  514. if context or out_of_date(code_path, img.filename(fmt),
  515. includes=code_includes):
  516. all_exists = False
  517. break
  518. img.formats.append(fmt)
  519. # assume that if we have one, we have them all
  520. if not all_exists:
  521. all_exists = (j > 0)
  522. break
  523. images.append(img)
  524. if not all_exists:
  525. break
  526. results.append((code_piece, images))
  527. else:
  528. all_exists = True
  529. if all_exists:
  530. return results
  531. # We didn't find the files, so build them
  532. results = []
  533. ns = plot_context if context else {}
  534. if context_reset:
  535. clear_state(config.plot_rcparams)
  536. plot_context.clear()
  537. close_figs = not context or close_figs
  538. for i, code_piece in enumerate(code_pieces):
  539. if not context or config.plot_apply_rcparams:
  540. clear_state(config.plot_rcparams, close_figs)
  541. elif close_figs:
  542. plt.close('all')
  543. _run_code(doctest.script_from_examples(code_piece) if is_doctest
  544. else code_piece,
  545. code_path, ns, function_name)
  546. images = []
  547. fig_managers = _pylab_helpers.Gcf.get_all_fig_managers()
  548. for j, figman in enumerate(fig_managers):
  549. if len(fig_managers) == 1 and len(code_pieces) == 1:
  550. img = ImageFile(output_base, output_dir)
  551. elif len(code_pieces) == 1:
  552. img = ImageFile("%s_%02d" % (output_base, j), output_dir)
  553. else:
  554. img = ImageFile("%s_%02d_%02d" % (output_base, i, j),
  555. output_dir)
  556. images.append(img)
  557. for fmt, dpi in formats:
  558. try:
  559. figman.canvas.figure.savefig(img.filename(fmt), dpi=dpi)
  560. if fmt == formats[0][0] and config.plot_srcset:
  561. # save a 2x, 3x etc version of the default...
  562. srcset = _parse_srcset(config.plot_srcset)
  563. for mult, suffix in srcset.items():
  564. fm = f'{suffix}.{fmt}'
  565. img.formats.append(fm)
  566. figman.canvas.figure.savefig(img.filename(fm),
  567. dpi=int(dpi * mult))
  568. except Exception as err:
  569. raise PlotError(traceback.format_exc()) from err
  570. img.formats.append(fmt)
  571. results.append((code_piece, images))
  572. if not context or config.plot_apply_rcparams:
  573. clear_state(config.plot_rcparams, close=not context)
  574. return results
  575. def run(arguments, content, options, state_machine, state, lineno):
  576. document = state_machine.document
  577. config = document.settings.env.config
  578. nofigs = 'nofigs' in options
  579. if config.plot_srcset and setup.app.builder.name == 'singlehtml':
  580. raise ExtensionError(
  581. 'plot_srcset option not compatible with single HTML writer')
  582. formats = get_plot_formats(config)
  583. default_fmt = formats[0][0]
  584. options.setdefault('include-source', config.plot_include_source)
  585. options.setdefault('show-source-link', config.plot_html_show_source_link)
  586. if 'class' in options:
  587. # classes are parsed into a list of string, and output by simply
  588. # printing the list, abusing the fact that RST guarantees to strip
  589. # non-conforming characters
  590. options['class'] = ['plot-directive'] + options['class']
  591. else:
  592. options.setdefault('class', ['plot-directive'])
  593. keep_context = 'context' in options
  594. context_opt = None if not keep_context else options['context']
  595. rst_file = document.attributes['source']
  596. rst_dir = os.path.dirname(rst_file)
  597. if len(arguments):
  598. if not config.plot_basedir:
  599. source_file_name = os.path.join(setup.app.builder.srcdir,
  600. directives.uri(arguments[0]))
  601. else:
  602. source_file_name = os.path.join(setup.confdir, config.plot_basedir,
  603. directives.uri(arguments[0]))
  604. # If there is content, it will be passed as a caption.
  605. caption = '\n'.join(content)
  606. # Enforce unambiguous use of captions.
  607. if "caption" in options:
  608. if caption:
  609. raise ValueError(
  610. 'Caption specified in both content and options.'
  611. ' Please remove ambiguity.'
  612. )
  613. # Use caption option
  614. caption = options["caption"]
  615. # If the optional function name is provided, use it
  616. if len(arguments) == 2:
  617. function_name = arguments[1]
  618. else:
  619. function_name = None
  620. code = Path(source_file_name).read_text(encoding='utf-8')
  621. output_base = os.path.basename(source_file_name)
  622. else:
  623. source_file_name = rst_file
  624. code = textwrap.dedent("\n".join(map(str, content)))
  625. counter = document.attributes.get('_plot_counter', 0) + 1
  626. document.attributes['_plot_counter'] = counter
  627. base, ext = os.path.splitext(os.path.basename(source_file_name))
  628. output_base = '%s-%d.py' % (base, counter)
  629. function_name = None
  630. caption = options.get('caption', '')
  631. base, source_ext = os.path.splitext(output_base)
  632. if source_ext in ('.py', '.rst', '.txt'):
  633. output_base = base
  634. else:
  635. source_ext = ''
  636. # ensure that LaTeX includegraphics doesn't choke in foo.bar.pdf filenames
  637. output_base = output_base.replace('.', '-')
  638. # is it in doctest format?
  639. is_doctest = contains_doctest(code)
  640. if 'format' in options:
  641. if options['format'] == 'python':
  642. is_doctest = False
  643. else:
  644. is_doctest = True
  645. # determine output directory name fragment
  646. source_rel_name = relpath(source_file_name, setup.confdir)
  647. source_rel_dir = os.path.dirname(source_rel_name).lstrip(os.path.sep)
  648. # build_dir: where to place output files (temporarily)
  649. build_dir = os.path.join(os.path.dirname(setup.app.doctreedir),
  650. 'plot_directive',
  651. source_rel_dir)
  652. # get rid of .. in paths, also changes pathsep
  653. # see note in Python docs for warning about symbolic links on Windows.
  654. # need to compare source and dest paths at end
  655. build_dir = os.path.normpath(build_dir)
  656. os.makedirs(build_dir, exist_ok=True)
  657. # how to link to files from the RST file
  658. try:
  659. build_dir_link = relpath(build_dir, rst_dir).replace(os.path.sep, '/')
  660. except ValueError:
  661. # on Windows, relpath raises ValueError when path and start are on
  662. # different mounts/drives
  663. build_dir_link = build_dir
  664. # get list of included rst files so that the output is updated when any
  665. # plots in the included files change. These attributes are modified by the
  666. # include directive (see the docutils.parsers.rst.directives.misc module).
  667. try:
  668. source_file_includes = [os.path.join(os.getcwd(), t[0])
  669. for t in state.document.include_log]
  670. except AttributeError:
  671. # the document.include_log attribute only exists in docutils >=0.17,
  672. # before that we need to inspect the state machine
  673. possible_sources = {os.path.join(setup.confdir, t[0])
  674. for t in state_machine.input_lines.items}
  675. source_file_includes = [f for f in possible_sources
  676. if os.path.isfile(f)]
  677. # remove the source file itself from the includes
  678. try:
  679. source_file_includes.remove(source_file_name)
  680. except ValueError:
  681. pass
  682. # save script (if necessary)
  683. if options['show-source-link']:
  684. Path(build_dir, output_base + source_ext).write_text(
  685. doctest.script_from_examples(code)
  686. if source_file_name == rst_file and is_doctest
  687. else code,
  688. encoding='utf-8')
  689. # make figures
  690. try:
  691. results = render_figures(code=code,
  692. code_path=source_file_name,
  693. output_dir=build_dir,
  694. output_base=output_base,
  695. context=keep_context,
  696. function_name=function_name,
  697. config=config,
  698. context_reset=context_opt == 'reset',
  699. close_figs=context_opt == 'close-figs',
  700. code_includes=source_file_includes)
  701. errors = []
  702. except PlotError as err:
  703. reporter = state.memo.reporter
  704. sm = reporter.system_message(
  705. 2, "Exception occurred in plotting {}\n from {}:\n{}".format(
  706. output_base, source_file_name, err),
  707. line=lineno)
  708. results = [(code, [])]
  709. errors = [sm]
  710. # Properly indent the caption
  711. if caption and config.plot_srcset:
  712. caption = f':caption: {caption}'
  713. elif caption:
  714. caption = '\n' + '\n'.join(' ' + line.strip()
  715. for line in caption.split('\n'))
  716. # generate output restructuredtext
  717. total_lines = []
  718. for j, (code_piece, images) in enumerate(results):
  719. if options['include-source']:
  720. if is_doctest:
  721. lines = ['', *code_piece.splitlines()]
  722. else:
  723. lines = ['.. code-block:: python', '',
  724. *textwrap.indent(code_piece, ' ').splitlines()]
  725. source_code = "\n".join(lines)
  726. else:
  727. source_code = ""
  728. if nofigs:
  729. images = []
  730. opts = [
  731. f':{key}: {val}' for key, val in options.items()
  732. if key in ('alt', 'height', 'width', 'scale', 'align', 'class')]
  733. # Not-None src_name signals the need for a source download in the
  734. # generated html
  735. if j == 0 and options['show-source-link']:
  736. src_name = output_base + source_ext
  737. else:
  738. src_name = None
  739. if config.plot_srcset:
  740. srcset = [*_parse_srcset(config.plot_srcset).values()]
  741. template = TEMPLATE_SRCSET
  742. else:
  743. srcset = None
  744. template = TEMPLATE
  745. result = jinja2.Template(config.plot_template or template).render(
  746. default_fmt=default_fmt,
  747. build_dir=build_dir_link,
  748. src_name=src_name,
  749. multi_image=len(images) > 1,
  750. options=opts,
  751. srcset=srcset,
  752. images=images,
  753. source_code=source_code,
  754. html_show_formats=config.plot_html_show_formats and len(images),
  755. caption=caption)
  756. total_lines.extend(result.split("\n"))
  757. total_lines.extend("\n")
  758. if total_lines:
  759. state_machine.insert_input(total_lines, source=source_file_name)
  760. return errors