This commit is contained in:
cjw
2026-02-12 23:22:11 +08:00
parent 7b09eb3d89
commit 89660bba4e
5988 changed files with 2517516 additions and 0 deletions
@@ -0,0 +1,736 @@
"""Plot directive module.
A directive for including a PyVista plot in a Sphinx document.
The ``.. pyvista-plot::`` sphinx directive will include an inline
``.png`` image.
The source code for the plot may be included in one of two ways:
1. Using **doctest** syntax::
.. pyvista-plot::
>>> import pyvista as pv
>>> sphere = pv.Sphere()
>>> out = sphere.plot()
2. **A path to a source file** as the argument to the directive::
.. pyvista-plot:: path/to/plot.py
When a path to a source file is given, the content of the
directive may optionally contain a caption for the plot::
.. pyvista-plot:: path/to/plot.py
The plot's caption.
Additionally, one may specify the name of a function to call (with
no arguments) immediately after importing the module::
.. pyvista-plot:: path/to/plot.py plot_function1
.. note::
Code blocks containing ``doctest:+SKIP`` will be skipped.
.. note::
Animations will not be saved, only the last frame will be shown.
**Options**
The ``pyvista-plot`` directive supports the following options:
include-source : bool
Whether to display the source code. The default can be changed
using the ``pyvista_plot_include_source`` variable in :file:`conf.py`.
encoding : str
If this source file is in a non-UTF8 or non-ASCII encoding, the
encoding must be specified using the ``:encoding:`` option. The
encoding will not be inferred using the ``-*- coding -*-`` metacomment.
context : None
If provided, the code will be run in the context of all previous plot
directives for which the ``:context:`` option was specified. This only
applies to inline code plot directives, not those run from files.
nofigs : None
When setting this flag, the code block will be run but no figures will be
inserted. This is usually useful with the ``:context:`` option.
caption : str
If specified, the option's argument will be used as a caption for the
figure. This overwrites the caption given in the content, when the plot
is generated from a file.
force_static : None
When setting this flag, static images will be used instead of an
interactive scene.
skip : bool, default: True
Whether to skip execution of this directive. If no argument is provided
i.e., ``:skip:``, then it defaults to ``:skip: true``. Default
behaviour is controlled by the ``plot_skip`` boolean variable in
:file:`conf.py`. Note that, if specified, this option overrides the
``plot_skip`` configuration.
optional : None
This flag marks the directive for *conditional* execution. Whether the
directive is executed is controlled by the ``plot_skip_optional``
boolean variable in :file:`conf.py`.
Additionally, this directive supports all the options of the `image`
directive, except for *target* (since plot will add its own target). These
include *alt*, *height*, *width*, *scale*, *align*.
**Configuration options**
.. versionchanged:: 0.45
Prior to v0.45, these directives conflicted with ``matplotlib``. All
directives have been prepended with ``pyvista_``.
The plot directive has the following configuration options:
pyvista_plot_include_source : bool, default: True
Default value for the ``include-source`` directive option.
Default is ``True``.
pyvista_plot_basedir : str
Base directory, to which ``plot::`` file names are relative
to. If ``None`` or unset, file names are relative to the
directory where the file containing the directive is.
pyvista_plot_html_show_formats : bool, default: True
Whether to show links to the files in HTML. Default ``True``.
pyvista_plot_template : str
Provide a customized Jinja2 template for preparing restructured text.
pyvista_plot_setup : str
Python code to be run before every plot directive block.
pyvista_plot_cleanup : str
Python code to be run after every plot directive block.
pyvista_plot_skip : bool, default: False
Default value for the ``skip`` directive option.
pyvista_plot_skip_optional : bool, default: False
Whether to skip execution of ``optional`` directives.
These options can be set by defining global variables of the same name in
:file:`conf.py`.
"""
from __future__ import annotations
import doctest
import os
from os.path import relpath
from pathlib import Path
import re
import shutil
import textwrap
import traceback
from typing import TYPE_CHECKING
from typing import ClassVar
from docutils.parsers.rst import Directive
from docutils.parsers.rst import directives
from docutils.parsers.rst.directives.images import Image
import jinja2 # Sphinx dependency.
# must enable BUILDING_GALLERY to keep windows active
# enable offscreen to hide figures when generating them.
import pyvista
if TYPE_CHECKING:
from collections.abc import Callable
pyvista.BUILDING_GALLERY = True
pyvista.OFF_SCREEN = True
# -----------------------------------------------------------------------------
# Registration hook
# -----------------------------------------------------------------------------
def _option_boolean(arg) -> bool:
if not arg or not arg.strip():
# no argument given, assume used as a flag
return True
elif arg.strip().lower() in ('no', '0', 'false'):
return False
elif arg.strip().lower() in ('yes', '1', 'true'):
return True
else: # pragma: no cover
msg = f'"{arg}" unknown boolean'
raise ValueError(msg)
def _option_context(arg):
if arg is not None: # pragma: no cover
msg = 'No arguments allowed for ``:context:``'
raise ValueError(msg)
def _option_format(arg):
return directives.choice(arg, ('python', 'doctest'))
class PlotDirective(Directive):
"""The ``.. pyvista-plot::`` directive, as documented in the module's docstring."""
has_content = True
required_arguments = 0
optional_arguments = 2
final_argument_whitespace = False
option_spec: ClassVar[dict[str, Callable]] = {
'alt': directives.unchanged,
'height': directives.length_or_unitless,
'width': directives.length_or_percentage_or_unitless,
'scale': directives.nonnegative_int,
'align': Image.align,
'include-source': _option_boolean,
'format': _option_format,
'context': _option_context,
'nofigs': directives.flag,
'encoding': directives.encoding,
'caption': directives.unchanged,
'force_static': directives.flag,
'skip': _option_boolean,
'optional': directives.flag,
}
def run(self):
"""Run the plot directive."""
try:
return run(
self.arguments,
self.content,
self.options,
self.state_machine,
self.state,
self.lineno,
)
except Exception as e: # noqa: BLE001 # pragma: no cover
raise self.error(str(e))
def setup(app):
"""Set up the plot directive."""
setup.app = app
setup.config = app.config
setup.confdir = app.confdir
app.add_directive('pyvista-plot', PlotDirective)
legacy_keys = [
'plot_include_source',
'plot_basedir',
'plot_html_show_formats',
'plot_template',
'plot_setup',
'plot_cleanup',
'plot_skip',
'plot_skip_optional',
]
def raise_on_legacy_config(app, config):
"""Raise a RuntimeError when using legacy configuration parameters.
These parameters conflict with matplotlib's ``plot_directive``.
"""
uses_matplotlib = 'matplotlib.sphinxext.plot_directive' in app.extensions
if not uses_matplotlib: # pragma: no cover
for key in legacy_keys:
if getattr(config, key, None) is not None:
msg = (
f"Sphinx config uses deprecated '{key}' without 'pyvista_' prefix. "
f"Rename it to 'pyvista_{key}"
)
raise RuntimeError(msg)
app.connect('config-inited', raise_on_legacy_config)
app.add_config_value('pyvista_plot_include_source', True, False)
app.add_config_value('pyvista_plot_basedir', None, True)
app.add_config_value('pyvista_plot_html_show_formats', True, True)
app.add_config_value('pyvista_plot_template', None, True)
app.add_config_value('pyvista_plot_setup', None, True)
app.add_config_value('pyvista_plot_cleanup', None, True)
app.add_config_value(name='pyvista_plot_skip', default=False, rebuild='html')
app.add_config_value(name='pyvista_plot_skip_optional', default=False, rebuild='html')
return {
'parallel_read_safe': True,
'parallel_write_safe': True,
'version': pyvista.__version__,
}
# -----------------------------------------------------------------------------
# Doctest handling
# -----------------------------------------------------------------------------
def _contains_doctest(text):
try:
# check if it's valid Python as-is
compile(text, '<string>', 'exec')
except SyntaxError:
pass
else:
return False
r = re.compile(r'^\s*>>>', re.MULTILINE)
m = r.search(text)
return bool(m)
def _contains_pyvista_plot(text) -> bool:
return '.. pyvista-plot::' in text
def _strip_comments(code):
"""Remove comments from a line of python code."""
return re.sub(r'(?m)^ *#.*\n?', '', code)
def _split_code_at_show(text):
"""Split code at plt.show() or plt.plot().
Includes logic to deal with edge cases like:
>>> import pyvista as pv
>>> pv.Sphere().plot(color='blue', cpos='xy')
>>> pv.Sphere().plot(color='red', cpos='xy')
"""
parts = []
is_doctest = _contains_doctest(text)
part = []
within_plot = False
for line in text.split('\n'):
part.append(line)
# check if show(...) or plot(...) is within the line
line_no_comments = _strip_comments(line)
if within_plot: # allow for multi-line plot(...
if line_no_comments.endswith(')'):
parts.append('\n'.join(part))
part = []
within_plot = False
elif _show_or_plot_in_string(line_no_comments):
if line_no_comments.endswith(')'):
parts.append('\n'.join(part))
part = []
else: # allow for multi-line plot(...
within_plot = True
if '\n'.join(part).strip():
parts.append('\n'.join(part))
return is_doctest, parts
def _show_or_plot_in_string(string):
# string contains `.show(`, `.plot(`, or `plot_xyz(` where `xyz` is one
# or more lower-case letters or underscore, e.g. `plot_cell(`, `plot_datasets(`
pattern = r'(?:\.plot\(|\.show\(|(?:[ \t\n.]plot_[a-z_]+?)\()'
return bool(re.search(pattern, string))
# -----------------------------------------------------------------------------
# Template
# -----------------------------------------------------------------------------
TEMPLATE = """
{{ source_code }}
.. only:: html
{% for img in images %}
{% if img.extension == 'vtksz' %}
.. tab-set::
.. tab-item:: Static Scene
.. figure:: {{ build_dir }}/{{ img.stem }}.png
{% for option in options -%}
{{ option }}
{% endfor %}
.. tab-item:: Interactive Scene
.. offlineviewer:: {{ build_dir }}/{{ img.stem }}.vtksz
{{ caption }} {# appropriate leading whitespace added beforehand #}
{% else %}
.. figure:: {{ build_dir }}/{{ img.basename }}
{% for option in options -%}
{{ option }}
{% endfor %}
{{ caption }} {# appropriate leading whitespace added beforehand #}
{% endif %}
{% endfor %}
"""
exception_template = """
.. only:: html
[`source code <%(linkdir)s/%(basename)s.py>`__]
Exception occurred rendering plot.
"""
# the context of the plot for all directives specified with the
# :context: option
plot_context = {}
class ImageFile:
"""Simple representation of an image file path."""
def __init__(self, dirname, basename):
"""Construct ImageFile."""
self.basename = basename
self.dirname = dirname
self.extension = Path(basename).suffix[1:]
@property
def filename(self):
"""Return the filename of this image."""
return str(Path(self.dirname) / self.basename)
@property
def stem(self):
"""Return the basename without the suffix."""
return Path(self.basename).stem
def __repr__(self) -> str: # pragma no cover
return self.filename
class PlotError(RuntimeError):
"""More descriptive plot error."""
def _run_code(*, code, code_path, ns=None, function_name=None): # noqa: ARG001
"""Run a docstring example.
Run the example if it does not contain ``'doctest:+SKIP'``, or a
```pyvista-plot::`` directive. In the later case, the doctest parser will
present the code-block again with the ```pyvista-plot::`` directive
and its options removed.
Import a Python module from a path, and run the function given by
name, if function_name is not None.
"""
# do not execute code containing any SKIP directives
if 'doctest:+SKIP' in code:
return ns
if 'pyvista-plot::' in code:
return ns
try:
if pyvista.PLOT_DIRECTIVE_THEME is not None:
pyvista.set_plot_theme(pyvista.PLOT_DIRECTIVE_THEME) # pragma: no cover
exec(code, ns)
except (Exception, SystemExit) as err: # pragma: no cover
raise PlotError(traceback.format_exc()) from err
return ns
def render_figures(
*,
code,
code_path,
output_dir,
output_base,
context,
function_name,
config,
force_static,
):
"""Run a pyplot script and save the images in *output_dir*.
Save the images under *output_dir* with file names derived from
*output_base*. Closed plotters are ignored if they were never
rendered.
"""
# We skip snippets that contain the ```pyvista-plot::`` directive as part of their code.
# The doctest parser will present the code-block once again with the ```pyvista-plot::``
# directive and its options properly parsed.
if _contains_pyvista_plot(code):
is_doctest = True
code_pieces = [code]
else:
# Try to determine if all images already exist
is_doctest, code_pieces = _split_code_at_show(code)
# Otherwise, we didn't find the files, so build them
results = []
ns = plot_context if context else {}
# Check for setup and teardown code for plots
code_setup = config.pyvista_plot_setup
code_cleanup = config.pyvista_plot_cleanup
if code_setup:
_run_code(code=code_setup, code_path=code_path, ns=ns, function_name=function_name)
try:
for i, code_piece in enumerate(code_pieces):
# generate the plot
_run_code(
code=doctest.script_from_examples(code_piece) if is_doctest else code_piece,
code_path=code_path,
ns=ns,
function_name=function_name,
)
images = []
if (
_show_or_plot_in_string(code_piece)
or '.open_gif' in code_piece
or 'plot=True' in code_piece
):
figures = pyvista.plotting.plotter._ALL_PLOTTERS
for j, (_, plotter) in enumerate(figures.items()):
if plotter._gif_filename is not None:
image_file = ImageFile(output_dir, f'{output_base}_{i:02d}_{j:02d}.gif')
shutil.move(plotter._gif_filename, image_file.filename)
else:
image_file = ImageFile(output_dir, f'{output_base}_{i:02d}_{j:02d}.png')
try:
plotter.screenshot(image_file.filename)
except RuntimeError: # pragma no cover
# ignore closed, unrendered plotters
continue
if force_static or (plotter.last_vtksz is None):
images.append(image_file)
continue
else:
image_file = ImageFile(
output_dir, f'{output_base}_{i:02d}_{j:02d}.vtksz'
)
with Path(image_file.filename).open('wb') as f:
f.write(plotter.last_vtksz)
images.append(image_file)
pyvista.close_all() # close and clear all plotters
results.append((code_piece, images))
finally:
if code_cleanup:
_run_code(code=code_cleanup, code_path=code_path, ns=ns, function_name=function_name)
return results
def run(arguments, content, options, state_machine, state, lineno): # noqa: PLR0917
"""Run the plot directive."""
document = state_machine.document
config = document.settings.env.config
nofigs = 'nofigs' in options
optional = 'optional' in options
force_static = 'force_static' in options
default_fmt = 'png'
options.setdefault('include-source', config.pyvista_plot_include_source)
options.setdefault('skip', config.pyvista_plot_skip)
skip = options['skip'] or (optional and config.pyvista_plot_skip_optional)
keep_context = 'context' in options
_ = None if not keep_context else options['context']
rst_file = document.attributes['source']
rst_dir = str(Path(rst_file).parent)
if len(arguments):
if not config.pyvista_plot_basedir:
source_file_name = str(Path(setup.app.builder.srcdir) / directives.uri(arguments[0]))
else:
source_file_name = str(
Path(setup.confdir) / config.pyvista_plot_basedir / directives.uri(arguments[0]),
)
# If there is content, it will be passed as a caption.
caption = '\n'.join(content)
# Enforce unambiguous use of captions.
if 'caption' in options:
if caption: # pragma: no cover
msg = 'Caption specified in both content and options. Please remove ambiguity.'
raise ValueError(msg)
# Use caption option
caption = options['caption']
# If the optional function name is provided, use it
function_name = arguments[1] if len(arguments) == 2 else None
code = Path(source_file_name).read_text(encoding='utf-8')
output_base = Path(source_file_name).name
else:
source_file_name = rst_file
code = textwrap.dedent('\n'.join(map(str, content)))
# note: this reuses the existing matplotlib plot counter if available
counter = document.attributes.get('_plot_counter', 0) + 1
document.attributes['_plot_counter'] = counter
base = Path(source_file_name).stem
ext = Path(source_file_name).suffix
output_base = f'{base}-{counter}{ext}'
function_name = None
caption = options.get('caption', '')
base = Path(output_base).stem
source_ext = Path(output_base).suffix
if source_ext in ('.py', '.rst', '.txt'):
output_base = base
else:
source_ext = ''
# ensure that LaTeX includegraphics doesn't choke in foo.bar.pdf filenames
output_base = output_base.replace('.', '-')
# is it in doctest format?
is_doctest = _contains_doctest(code)
if 'format' in options:
is_doctest = options['format'] != 'python'
# determine output directory name fragment
source_rel_name = relpath(source_file_name, setup.confdir)
source_rel_dir = str(Path(source_rel_name).parent).lstrip(os.path.sep)
# build_dir: where to place output files (temporarily)
build_dir = str(Path(setup.app.doctreedir).parent / 'plot_directive' / source_rel_dir)
# get rid of .. in paths, also changes pathsep
# see note in Python docs for warning about symbolic links on Windows.
# need to compare source and dest paths at end
build_dir = os.path.normpath(build_dir)
Path(build_dir).mkdir(parents=True, exist_ok=True)
# output_dir: final location in the builder's directory
dest_dir = str((Path(setup.app.builder.outdir) / source_rel_dir).resolve())
Path(dest_dir).mkdir(parents=True, exist_ok=True)
# how to link to files from the RST file
dest_dir_link = os.path.join( # noqa: PTH118
relpath(setup.confdir, rst_dir),
source_rel_dir,
).replace(os.path.sep, '/')
try:
build_dir_link = relpath(build_dir, rst_dir).replace(os.path.sep, '/')
except ValueError: # pragma: no cover
# on Windows, relpath raises ValueError when path and start are on
# different mounts/drives
build_dir_link = build_dir
_ = dest_dir_link + '/' + output_base + source_ext
# make figures
errors = []
if skip:
results = [(code, [])]
else:
try:
results = render_figures(
code=code,
code_path=source_file_name,
output_dir=build_dir,
output_base=output_base,
context=keep_context,
function_name=function_name,
config=config,
force_static=force_static,
)
except PlotError as err: # pragma: no cover
reporter = state.memo.reporter
sm = reporter.system_message(
2,
f'Exception occurred in plotting {output_base}\n from {source_file_name}:\n{err}',
line=lineno,
)
results = [(code, [])]
errors.append([sm])
# Properly indent the caption
caption = (
'' if skip else '\n' + '\n'.join(' ' + line.strip() for line in caption.split('\n'))
)
# generate output restructuredtext
total_lines = []
for _, (code_piece, images) in enumerate(results):
if options['include-source']:
if is_doctest:
lines = ['', *code_piece.splitlines()]
else:
lines = [
'.. code-block:: python',
'',
*textwrap.indent(code_piece, ' ').splitlines(),
]
source_code = '\n'.join(lines)
else:
source_code = ''
images_input = [] if nofigs else images
opts = [
f':{key}: {val}'
for key, val in options.items()
if key in ('alt', 'height', 'width', 'scale', 'align')
]
result = jinja2.Template(config.pyvista_plot_template or TEMPLATE).render(
default_fmt=default_fmt,
dest_dir=dest_dir_link,
build_dir=build_dir_link,
source_link=None,
multi_image=len(images_input) > 1,
options=opts,
images=images_input,
source_code=source_code,
html_show_formats=config.pyvista_plot_html_show_formats and len(images_input),
caption=caption,
)
total_lines.extend(result.split('\n'))
total_lines.extend('\n')
if total_lines:
state_machine.insert_input(total_lines, source=source_file_name)
# copy image files to builder's output directory, if necessary
Path(dest_dir).mkdir(parents=True, exist_ok=True)
for _, images in results:
for image in images:
destimg = str(Path(dest_dir) / image.basename)
if image.filename != destimg:
shutil.copyfile(image.filename, destimg)
# copy script (if necessary)
Path(dest_dir, output_base + source_ext).write_text(
doctest.script_from_examples(code)
if source_file_name == rst_file and is_doctest
else code,
encoding='utf-8',
)
return errors