mirror of
https://github.com/aykhans/AzSuicideDataVisualization.git
synced 2025-04-22 10:28:02 +00:00
342 lines
11 KiB
Python
342 lines
11 KiB
Python
"""
|
|
Altair Plot Sphinx Extension
|
|
============================
|
|
|
|
This extension provides a means of inserting live-rendered Altair plots within
|
|
sphinx documentation. There are two directives defined: ``altair-setup`` and
|
|
``altair-plot``. ``altair-setup`` code is used to set-up various options
|
|
prior to running the plot code. For example::
|
|
|
|
.. altair-plot::
|
|
:output: none
|
|
|
|
from altair import *
|
|
import pandas as pd
|
|
data = pd.DataFrame({'a': list('CCCDDDEEE'),
|
|
'b': [2, 7, 4, 1, 2, 6, 8, 4, 7]})
|
|
|
|
.. altair-plot::
|
|
|
|
Chart(data).mark_point().encode(
|
|
x='a',
|
|
y='b'
|
|
)
|
|
|
|
In the case of the ``altair-plot`` code, the *last statement* of the code-block
|
|
should contain the chart object you wish to be rendered.
|
|
|
|
Options
|
|
-------
|
|
The directives have the following options::
|
|
|
|
.. altair-plot::
|
|
:namespace: # specify a plotting namespace that is persistent within the doc
|
|
:hide-code: # if set, then hide the code and only show the plot
|
|
:code-below: # if set, then code is below rather than above the figure
|
|
:output: [plot|repr|stdout|none]
|
|
:alt: text # Alternate text when plot cannot be rendered
|
|
:links: editor source export # specify one or more of these options
|
|
:chart-var-name: chart # name of variable in namespace containing output
|
|
|
|
|
|
Additionally, this extension introduces a global configuration
|
|
``altairplot_links``, set in your ``conf.py`` which is a dictionary
|
|
of links that will appear below plots, unless the ``:links:`` option
|
|
again overrides it. It should look something like this::
|
|
|
|
# conf.py
|
|
# ...
|
|
altairplot_links = {'editor': True, 'source': True, 'export': True}
|
|
# ...
|
|
|
|
If this configuration is not specified, all are set to True.
|
|
"""
|
|
|
|
import contextlib
|
|
import io
|
|
import os
|
|
import json
|
|
import warnings
|
|
|
|
import jinja2
|
|
|
|
from docutils import nodes
|
|
from docutils.parsers.rst import Directive
|
|
from docutils.parsers.rst.directives import flag, unchanged
|
|
|
|
from sphinx.locale import _
|
|
|
|
import altair as alt
|
|
from altair.utils.execeval import eval_block
|
|
|
|
# These default URLs can be changed in conf.py; see setup() below.
|
|
VEGA_JS_URL_DEFAULT = "https://cdn.jsdelivr.net/npm/vega@{}".format(alt.VEGA_VERSION)
|
|
VEGALITE_JS_URL_DEFAULT = "https://cdn.jsdelivr.net/npm/vega-lite@{}".format(
|
|
alt.VEGALITE_VERSION
|
|
)
|
|
VEGAEMBED_JS_URL_DEFAULT = "https://cdn.jsdelivr.net/npm/vega-embed@{}".format(
|
|
alt.VEGAEMBED_VERSION
|
|
)
|
|
|
|
|
|
VGL_TEMPLATE = jinja2.Template(
|
|
"""
|
|
<div id="{{ div_id }}">
|
|
<script>
|
|
// embed when document is loaded, to ensure vega library is available
|
|
// this works on all modern browsers, except IE8 and older
|
|
document.addEventListener("DOMContentLoaded", function(event) {
|
|
var spec = {{ spec }};
|
|
var opt = {
|
|
"mode": "{{ mode }}",
|
|
"renderer": "{{ renderer }}",
|
|
"actions": {{ actions}}
|
|
};
|
|
vegaEmbed('#{{ div_id }}', spec, opt).catch(console.err);
|
|
});
|
|
</script>
|
|
</div>
|
|
"""
|
|
)
|
|
|
|
|
|
class altair_plot(nodes.General, nodes.Element):
|
|
pass
|
|
|
|
|
|
def purge_altair_namespaces(app, env, docname):
|
|
if not hasattr(env, "_altair_namespaces"):
|
|
return
|
|
env._altair_namespaces.pop(docname, {})
|
|
|
|
|
|
DEFAULT_ALTAIRPLOT_LINKS = {"editor": True, "source": True, "export": True}
|
|
|
|
|
|
def validate_links(links):
|
|
if links.strip().lower() == "none":
|
|
return False
|
|
|
|
links = links.strip().split()
|
|
diff = set(links) - set(DEFAULT_ALTAIRPLOT_LINKS.keys())
|
|
if diff:
|
|
raise ValueError("Following links are invalid: {}".format(list(diff)))
|
|
return {link: link in links for link in DEFAULT_ALTAIRPLOT_LINKS}
|
|
|
|
|
|
def validate_output(output):
|
|
output = output.strip().lower()
|
|
if output not in ["plot", "repr", "stdout", "none"]:
|
|
raise ValueError(":output: flag must be one of [plot|repr|stdout|none]")
|
|
return output
|
|
|
|
|
|
class AltairPlotDirective(Directive):
|
|
has_content = True
|
|
|
|
option_spec = {
|
|
"hide-code": flag,
|
|
"code-below": flag,
|
|
"namespace": unchanged,
|
|
"output": validate_output,
|
|
"alt": unchanged,
|
|
"links": validate_links,
|
|
"chart-var-name": unchanged,
|
|
}
|
|
|
|
def run(self):
|
|
env = self.state.document.settings.env
|
|
app = env.app
|
|
|
|
show_code = "hide-code" not in self.options
|
|
code_below = "code-below" in self.options
|
|
|
|
if not hasattr(env, "_altair_namespaces"):
|
|
env._altair_namespaces = {}
|
|
namespace_id = self.options.get("namespace", "default")
|
|
namespace = env._altair_namespaces.setdefault(env.docname, {}).setdefault(
|
|
namespace_id, {}
|
|
)
|
|
|
|
code = "\n".join(self.content)
|
|
|
|
if show_code:
|
|
source_literal = nodes.literal_block(code, code)
|
|
source_literal["language"] = "python"
|
|
|
|
# get the name of the source file we are currently processing
|
|
rst_source = self.state_machine.document["source"]
|
|
rst_dir = os.path.dirname(rst_source)
|
|
rst_filename = os.path.basename(rst_source)
|
|
|
|
# use the source file name to construct a friendly target_id
|
|
serialno = env.new_serialno("altair-plot")
|
|
rst_base = rst_filename.replace(".", "-")
|
|
div_id = "{}-altair-plot-{}".format(rst_base, serialno)
|
|
target_id = "{}-altair-source-{}".format(rst_base, serialno)
|
|
target_node = nodes.target("", "", ids=[target_id])
|
|
|
|
# create the node in which the plot will appear;
|
|
# this will be processed by html_visit_altair_plot
|
|
plot_node = altair_plot()
|
|
plot_node["target_id"] = target_id
|
|
plot_node["div_id"] = div_id
|
|
plot_node["code"] = code
|
|
plot_node["namespace"] = namespace
|
|
plot_node["relpath"] = os.path.relpath(rst_dir, env.srcdir)
|
|
plot_node["rst_source"] = rst_source
|
|
plot_node["rst_lineno"] = self.lineno
|
|
plot_node["links"] = self.options.get(
|
|
"links", app.builder.config.altairplot_links
|
|
)
|
|
plot_node["output"] = self.options.get("output", "plot")
|
|
plot_node["chart-var-name"] = self.options.get("chart-var-name", None)
|
|
|
|
if "alt" in self.options:
|
|
plot_node["alt"] = self.options["alt"]
|
|
|
|
result = [target_node]
|
|
|
|
if code_below:
|
|
result += [plot_node]
|
|
if show_code:
|
|
result += [source_literal]
|
|
if not code_below:
|
|
result += [plot_node]
|
|
|
|
return result
|
|
|
|
|
|
def html_visit_altair_plot(self, node):
|
|
# Execute the code, saving output and namespace
|
|
namespace = node["namespace"]
|
|
try:
|
|
f = io.StringIO()
|
|
with contextlib.redirect_stdout(f):
|
|
chart = eval_block(node["code"], namespace)
|
|
stdout = f.getvalue()
|
|
except Exception as e:
|
|
warnings.warn(
|
|
"altair-plot: {}:{} Code Execution failed:"
|
|
"{}: {}".format(
|
|
node["rst_source"], node["rst_lineno"], e.__class__.__name__, str(e)
|
|
)
|
|
)
|
|
raise nodes.SkipNode
|
|
|
|
chart_name = node["chart-var-name"]
|
|
if chart_name is not None:
|
|
if chart_name not in namespace:
|
|
raise ValueError(
|
|
"chart-var-name='{}' not present in namespace" "".format(chart_name)
|
|
)
|
|
chart = namespace[chart_name]
|
|
|
|
output = node["output"]
|
|
|
|
if output == "none":
|
|
raise nodes.SkipNode
|
|
elif output == "stdout":
|
|
if not stdout:
|
|
raise nodes.SkipNode
|
|
else:
|
|
output_literal = nodes.literal_block(stdout, stdout)
|
|
output_literal["language"] = "none"
|
|
node.extend([output_literal])
|
|
elif output == "repr":
|
|
if chart is None:
|
|
raise nodes.SkipNode
|
|
else:
|
|
rep = " " + repr(chart).replace("\n", "\n ")
|
|
repr_literal = nodes.literal_block(rep, rep)
|
|
repr_literal["language"] = "none"
|
|
node.extend([repr_literal])
|
|
elif output == "plot":
|
|
if isinstance(chart, alt.TopLevelMixin):
|
|
# Last line should be a chart; convert to spec dict
|
|
try:
|
|
spec = chart.to_dict()
|
|
except alt.utils.schemapi.SchemaValidationError:
|
|
raise ValueError("Invalid chart: {0}".format(node["code"]))
|
|
actions = node["links"]
|
|
|
|
# TODO: add an option to save spects to file & load from there.
|
|
# TODO: add renderer option
|
|
|
|
# Write spec to a *.vl.json file
|
|
# dest_dir = os.path.join(self.builder.outdir, node['relpath'])
|
|
# if not os.path.exists(dest_dir):
|
|
# os.makedirs(dest_dir)
|
|
# filename = "{0}.vl.json".format(node['target_id'])
|
|
# dest_path = os.path.join(dest_dir, filename)
|
|
# with open(dest_path, 'w') as f:
|
|
# json.dump(spec, f)
|
|
|
|
# Pass relevant info into the template and append to the output
|
|
html = VGL_TEMPLATE.render(
|
|
div_id=node["div_id"],
|
|
spec=json.dumps(spec),
|
|
mode="vega-lite",
|
|
renderer="canvas",
|
|
actions=json.dumps(actions),
|
|
)
|
|
self.body.append(html)
|
|
else:
|
|
warnings.warn(
|
|
"altair-plot: {}:{} Malformed block. Last line of "
|
|
"code block should define a valid altair Chart object."
|
|
"".format(node["rst_source"], node["rst_lineno"])
|
|
)
|
|
raise nodes.SkipNode
|
|
|
|
|
|
def generic_visit_altair_plot(self, node):
|
|
# TODO: generate PNGs and insert them here
|
|
if "alt" in node.attributes:
|
|
self.body.append(_("[ graph: %s ]") % node["alt"])
|
|
else:
|
|
self.body.append(_("[ graph ]"))
|
|
raise nodes.SkipNode
|
|
|
|
|
|
def depart_altair_plot(self, node):
|
|
return
|
|
|
|
|
|
def builder_inited(app):
|
|
app.add_js_file(app.config.altairplot_vega_js_url)
|
|
app.add_js_file(app.config.altairplot_vegalite_js_url)
|
|
app.add_js_file(app.config.altairplot_vegaembed_js_url)
|
|
|
|
|
|
def setup(app):
|
|
setup.app = app
|
|
setup.config = app.config
|
|
setup.confdir = app.confdir
|
|
|
|
app.add_config_value("altairplot_links", DEFAULT_ALTAIRPLOT_LINKS, "env")
|
|
|
|
app.add_config_value("altairplot_vega_js_url", VEGA_JS_URL_DEFAULT, "html")
|
|
app.add_config_value("altairplot_vegalite_js_url", VEGALITE_JS_URL_DEFAULT, "html")
|
|
app.add_config_value(
|
|
"altairplot_vegaembed_js_url", VEGAEMBED_JS_URL_DEFAULT, "html"
|
|
)
|
|
|
|
app.add_directive("altair-plot", AltairPlotDirective)
|
|
|
|
app.add_css_file("altair-plot.css")
|
|
|
|
app.add_node(
|
|
altair_plot,
|
|
html=(html_visit_altair_plot, depart_altair_plot),
|
|
latex=(generic_visit_altair_plot, depart_altair_plot),
|
|
texinfo=(generic_visit_altair_plot, depart_altair_plot),
|
|
text=(generic_visit_altair_plot, depart_altair_plot),
|
|
man=(generic_visit_altair_plot, depart_altair_plot),
|
|
)
|
|
|
|
app.connect("env-purge-doc", purge_altair_namespaces)
|
|
app.connect("builder-inited", builder_inited)
|
|
|
|
return {"version": "0.1"}
|