Repository: altair-viz/altair
Pull request #213 (2016-09-24 13:42:40)

MAINT: improve nodejs interaction for savechart()


Patch:
diff --git a/altair/api.py b/altair/api.py
index 78e4cb2b..a7f34453 100644
--- a/altair/api.py
+++ b/altair/api.py
@@ -3,11 +3,14 @@ Main API for Vega-lite spec generation.
 
 DSL mapping Vega types to IPython traitlets.
 """
+import os
+
 import traitlets as T
 import pandas as pd
 
 from .utils import visitors
 from .utils._py3k_compat import string_types
+from .utils import node
 
 from . import schema
 
@@ -67,6 +70,66 @@ def use_signature(Obj):
 # Top-level Objects
 #*************************************************************************
 class TopLevelMixin(object):
+    @staticmethod
+    def _png_output_available():
+        return node.vl_cmd_available('vl2png')
+
+    @staticmethod
+    def _svg_output_available():
+        return node.vl_cmd_available('vl2svg')
+
+    def savechart(self, outfile, filetype=None):
+        """Save a chart to file, in either png, svg, json, or html format.
+
+        Note that png/svg output requires several nodejs packages to be
+        installed and correctly configured. Before running this, you must
+        have nodejs and cairo on your system and use the node package manager
+        to install the ``canvas`` and ``vega-lite`` packages.
+
+        If you are using anaconda, you can set it up this way:
+
+            $ conda create -n node-env -c conda-forge python=2.7 cairo nodejs altair
+            $ source activate node-env
+            $ npm install canvas vega-lite
+
+        The node binaries used here (``vl2vg``, ``vl2png``, ``vl2svg``) will be
+        installed in the node root directory, which should be automatically
+        detected by this function.
+
+        Parameters
+        ----------
+        outfile : str
+            The output filename
+        filetype : str (optional)
+            The filetype to use. One of ('svg', 'png', 'json', 'html').
+            If not specified, it will be inferred from outfile.
+        """
+        if filetype is None:
+            try:
+                base, ext = os.path.splitext(outfile)
+                filetype = ext[1:]
+            except AttributeError:
+                raise ValueError('filetype could not be inferred')
+
+        if filetype in node.SUPPORTED_FILETYPES:
+            node.savechart(self, outfile, filetype)
+        elif filetype == 'json':
+            if hasattr(outfile, 'write'):
+                outfile.write(self.to_json())
+            else:
+                with open(outfile, 'w') as f:
+                    f.write(self.to_json())
+        elif filetype == 'html':
+            if hasattr(outfile, 'write'):
+                outfile.write(self.to_html())
+            else:
+                with open(outfile, 'w') as f:
+                    f.write(self.to_html())
+        else:
+            supported = _node.SUPPORTED_FILETYPES + ['json', 'html']
+            raise ValueError('Cannot save chart of type {0}; supported'
+                             'extensions are {1}'.format(filetype, supported))
+
     def to_html(self, template=None, title=None, **kwargs):
         """Emit a stand-alone HTML document containing this chart.
 
diff --git a/altair/utils/node.py b/altair/utils/node.py
index c1033d1c..81107c9f 100644
--- a/altair/utils/node.py
+++ b/altair/utils/node.py
@@ -1,129 +1,114 @@
-"""
-Utilities for using nodejs packages for generating png & svg
-"""
+"""Commands to interact with vega/vega-lite nodejs executables"""
 from __future__ import print_function
 
 import os
-import tempfile
-import subprocess
+import json
+from contextlib import contextmanager
+from xml.etree import ElementTree
+from subprocess import Popen, PIPE, check_output, CalledProcessError
 
 
-__all__ = ['savechart', 'savechart_available', 'NodeExecError']
-
+COMMANDS = ['vl2vg', 'vl2png', 'vl2svg']
 SUPPORTED_FILETYPES = ['png', 'svg']
+TEST_SPEC = {'data': {'values': [{'x': 1, 'y': 1}, {'x': 2, 'y': 2}]},
+             'encoding': {'x': {'field': 'x', 'type': 'quantitative'},
+                          'y': {'field': 'y', 'type': 'quantitative'}},
+             'mark': 'point'}
+
+
+def consistent_with_png(filename):
+    """Check that the file has a header consistent with PNG"""
+    with open(filename, 'rb') as f:
+        arr = f.read()
+    return arr.startswith(b'\x89PNG\r\n')
+
+
+def consistent_with_svg(filename):
+    """Check that the file has a structure consistent with SVG"""
+    try:
+        e = ElementTree.parse(filename)
+    except ElementTree.ParseError:
+        return False
+    else:
+        keys = e.getroot().keys()
+        return set(keys) == {'width', 'version', 'class', 'height'}
+
+
+@contextmanager
+def ensure_npm_bin_in_path(verbose=False):
+    if verbose:
+        print('> npm bin')
+    try:
+        npm_bin = check_output(['npm', 'bin'])
+    except OSError:
+        yield
+    else:
+        if hasattr(npm_bin, 'decode'):
+            npm_bin = npm_bin.decode('utf-8')
+        npm_bin = npm_bin.strip()
+
+        old_path = os.environ['PATH']
+
+        if npm_bin not in os.environ['PATH'].split(os.pathsep):
+            os.environ["PATH"] += os.pathsep + npm_bin
+        yield
+        os.environ["PATH"] = old_path
+
+
+def vl_cmd_available(cmd, verbose=False):
+    spec = json.dumps(TEST_SPEC)
+    with ensure_npm_bin_in_path(verbose):
+        try:
+            if verbose:
+                print('> ' + ' '.join(cmd))
+            p = Popen([cmd], stdout=PIPE, stdin=PIPE, stderr=PIPE)
+        except OSError:
+            return False
 
+        out, err = p.communicate(input='{0}'.format(spec).encode())
+    return not p.returncode
 
-class NodeExecError(ValueError):
-    pass
 
+def _convert_vegalite_spec(spec, cmd, outfile=None, verbose=False):
+    if verbose:
+        print('> ' + ' '.join(cmd))
+    p = Popen([cmd], stdout=PIPE, stdin=PIPE, stderr=PIPE)
 
-class _NodeExecutor(object):
-    """Class to execute vega/vega-lite conversion commands in nodejs
-
-    Parameters
-    ----------
-    node_bin_dir : str
-        The directory containing binary executables installed by node.
-        If not specified, then ``npm root`` will be used to find it.
-    verbose : bool (optional)
-        If True (default) then print commands before executing them.
-    """
-    def __init__(self, node_bin_dir=None, verbose=True):
-        self.verbose = verbose
-        self.node_bin_dir = os.path.abspath(node_bin_dir or self.npm_root)
-
-    def _exec(self, executable, inputfile, outputfile):
-        full_executable = os.path.join(self.node_bin_dir, executable)
-        if not os.path.exists(full_executable):
-            raise ValueError('{0} not found'.format(full_executable))
-        sp = subprocess.Popen([full_executable, inputfile],
-                              stdout=subprocess.PIPE,
-                              stderr=subprocess.PIPE)
-        out, err = sp.communicate()
-        if sp.returncode:
-            raise NodeExecError('``{0}`` failed: {1}'
-                                ''.format(executable, err))
-        with open(outputfile, 'wb') as f:
-            f.write(out)
-
-    def _chain(self, executable1, executable2, inputfile, outputfile):
-        # stdin/stdout piping is not yet supported
-        # see https://github.com/vega/vega/issues/612
-        # We'll use temporary files instead.
-        with tempfile.NamedTemporaryFile(suffix='.json') as f:
-            self._exec(executable1, inputfile, f.name)
-            self._exec(executable2, f.name, outputfile)
-
-    def vl2vg(self, inputfile, outputfile):
-        return _exec('vl2vg', inputfile, outputfile)
-
-    def vg2png(self, inputfile, outputfile):
-        return _exec('vl2vg', inputfile, outputfile)
-
-    def vg2svg(self, inputfile, outputfile):
-        return _exec('vl2vg', inputfile, outputfile)
-
-    def vl2png(self, inputfile, outputfile):
-        self._chain('vl2vg', 'vg2png', inputfile, outputfile)
-
-    def vl2svg(self, inputfile, outputfile):
-        self._chain('vl2vg', 'vg2svg', inputfile, outputfile)
-
-    @property
-    def npm_root(self):
-        sp = subprocess.Popen(['npm', 'root'],
-                              stdout=subprocess.PIPE,
-                              stderr=subprocess.PIPE)
-        rootdir, err = sp.communicate()
-        if sp.returncode:
-            raise NodeExecError("``npm root`` failed: {0}".format(err))
-
-        if hasattr(rootdir, 'decode'):
-            # decode bytes in Pytho 3
-            rootdir = rootdir.strip().decode('utf-8')
-        return os.path.join(rootdir, '.bin')
-
-    def savechart_available(self):
-        try:
-            root = self.npm_root
-        except:
-            return False
-        else:
-            root = root.strip()
-            if not os.path.exists(root):
-                return False
-            else:
-                L = os.listdir(root)
-                return ('vl2vg' in L and 'vg2png' in L and 'vg2svg' in L)
+    input_ = '{0}'.format(json.dumps(spec))
+    if hasattr(input_, 'encode'):
+        input_ = input_.encode()
 
+    out, err = p.communicate(input=input_)
+    if p.returncode:
+        raise CalledProcessError(p.returncode, cmd, output=out, stderr=err)
 
-def savechart_available(node_bin_dir=None):
-    """Check if savechart() command is available on your system"""
-    node = _NodeExecutor(node_bin_dir=node_bin_dir)
-    return node.savechart_available()
-        
+    if outfile is not None:
+        if hasattr(outfile, 'write'):
+            outfile.write(out)
+        else:
+            with open(outfile, 'wb') as f:
+                f.write(out)
+    else:
+        return out
 
 
-def savechart(chart, filename, filetype=None,
-              node_bin_dir=None, verbose=False):
-    """**EXPERIMENTAL** function to save a chart to png or svg
+def savechart(chart, filename, filetype=None, verbose=False):
+    """Save a chart to png or svg
 
     Note that this requires several nodejs packages to be installed and
     correctly configured. Before running this, you must have nodejs on
     your system and use the node package manager and install the ``canvas``
     and ``vega-lite`` packages.
 
-    If you are using anaconda, you can set it up this way:
+    If you are using anaconda, you can set it up this way::
 
-        $ conda create -n node-env -c conda-forge python=2.7 nodejs altair
+        $ conda create -n node-env -c conda-forge python=2.7 cairo nodejs altair
         $ source activate node-env
         $ npm install canvas vega-lite
 
-    The node binaries used here (``vl2vg``, ``vg2png``, ``vg2svg``) will be
+    The node binaries used here (``vl2vg``, ``vl2png``, ``vl2svg``) will be
     installed in the node root directory, which should be automatically
-    detected by this function. If you have these nodejs packages installed
-    and this function doesn't work, try explicitly passing their path using
-    the ``node_bin_dir`` argument.
+    detected by this function.
 
     Parameters
     ----------
@@ -134,9 +119,6 @@ def savechart(chart, filename, filetype=None,
     filetype : str (optional)
         The filetype to use (either 'svg' or 'png'). If not specified,
         it will be inferred from the filename.
-    node_bin_dir : str
-        The directory containing binary executables installed by node.
-        If not specified, then ``npm root`` will be used to determine it.
     verbose : bool (optional)
         If True (default) then print commands before executing them.
     """
@@ -148,17 +130,8 @@ def savechart(chart, filename, filetype=None,
         raise ValueError("Filetype {0} not valid: must be one of {1}"
                          "".format(filetype, SUPPORTED_FILETYPES))
 
-    node = _NodeExecutor(node_bin_dir=node_bin_dir,
-                         verbose=verbose)
-    if not node.savechart_available():
-        raise ValueError("Must install and configure nodejs tools "
-                         "to use savechart")
-        
-    
-    # stdin/stdout piping is not yet supported
-    # see https://github.com/vega/vega/issues/612
-    # We'll use temporary files instead.
-    with tempfile.NamedTemporaryFile(suffix='.json') as f:
-        with open(f.name, 'w') as ff:
-            ff.write(chart.to_json())
-        getattr(node, 'vl2' + filetype)(f.name, filename)
+    with ensure_npm_bin_in_path(verbose=verbose):
+        _convert_vegalite_spec(chart.to_dict(),
+                               'vl2' + filetype,
+                               outfile=filename,
+                               verbose=verbose)
diff --git a/doc/conf.py b/doc/conf.py
index 4969fa08..b8c8204b 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -139,7 +139,7 @@ html_theme = 'haiku'
 #html_title = 'altair v1.0.0'
 
 # A shorter title for the navigation bar.  Default is the same as html_title.
-#html_short_title = None
+html_short_title = 'Altair'
 
 # The name of an image file (relative to this directory) to place at the top
 # of the sidebar.
diff --git a/doc/documentation/displaying.rst b/doc/documentation/displaying.rst
index 8a967ffb..57474fea 100644
--- a/doc/documentation/displaying.rst
+++ b/doc/documentation/displaying.rst
@@ -141,14 +141,15 @@ If you don't have ``nodejs`` and are using ``conda``, you can create an
 environment with nodejs/npm and the required packages as follows
 (note that the canvas tool seems to require Python 2.7):
 
-    $ conda create -n nodejs-env -c conda-forge python=2.7 nodejs altair
+    $ conda create -n nodejs-env -c conda-forge python=2.7 nodejs cairo altair
     $ source activate nodejs-env
     $ npm install canvas vega-lite
 
 Once you have successfully installed these packages, you should have new binary
 files ``vl2vg``, ``vg2png``, and ``vg2eps`` within your node root directory.
 
-With this set-up, you can use the installed command-line utilities to save your chart:
+With these packages properly installed, you can use the ``savechart`` method
+to save a chart to file:
 
 >>> from altair import Chart, load_dataset
 >>> data = load_dataset('cars', url_only=True)
@@ -157,20 +158,16 @@ With this set-up, you can use the installed command-line utilities to save your
 ...             y='Miles_per_Gallon:Q',
 ...             color='Origin:N',
 ...         )
-
->>> from altair.utils.node import savechart, savechart_available
 >>> # save as PNG
->>> savechart(chart, 'mychart.png')  # doctest: +SKIP
+>>> chart.savechart('mychart.png')  # doctest: +SKIP
 >>> # save as SVG
->>> savechart(chart, 'mychart.svg')  # doctest: +SKIP
-
-Internally, the code executes ``npm root`` to determine the installation directory,
-and calls the ``vl2vg``, ``vg2png``, and ``vg2eps`` executables within this root.
-If you want to manually specify the npm executable directory, you can use the
-``node_bin_dir`` keyword:
+>>> chart.savechart('mychart.svg')  # doctest: +SKIP
 
->>> savechart(chart, 'mychart.svg', node_bin_dir='~/node_modules/.bin/')
+Internally, this requires the ``vl2png`` or ``vl2svg`` executables, which
+must either be in the system ``$PATH`` variable, or within the node
+binary directory specified by the command ``npm bin``.
 
+This extra installation step is straightforward, but admittedly a bit clunky.
 We hope to find a way to streamline this in the future, but creating transparent
 interactions between Python packages and NodeJS packages remains challenging.
 If you have ideas on how to improve this aspect of Altair's user experience,


Test patch:
diff --git a/altair/tests/test_api.py b/altair/tests/test_api.py
index 686bdf08..0d23c772 100644
--- a/altair/tests/test_api.py
+++ b/altair/tests/test_api.py
@@ -1,6 +1,7 @@
 import pytest
 import warnings
 import json
+import tempfile
 
 import numpy as np
 import pandas as pd
@@ -9,9 +10,15 @@ from .. import *
 from .. import schema
 from ..examples import iter_examples
 from ..utils import parse_shorthand, infer_vegalite_type
+from ..utils.node import consistent_with_png, consistent_with_svg
 from ..datasets import connection_ok
 from ..utils._py3k_compat import PY2
 
+def make_chart():
+    data = pd.DataFrame({'x': range(10),
+                         'y': range(10)})
+    return Chart(data).mark_point().encode(x='x', y='y')
+
 
 def test_chart_url_input():
     url = 'http://vega.github.io/vega-lite/data/'
@@ -40,6 +47,45 @@ def test_chart_to_json_round_trip():
     assert chart1.to_dict() == chart2.to_dict()
 
 
+@pytest.mark.skipif(not Chart._png_output_available(),
+                    reason='command-line tool vl2png is not available')
+def test_savechart_png():
+    chart = make_chart()
+
+    with tempfile.NamedTemporaryFile(suffix='.png') as f:
+        chart.savechart(f.name)
+        assert consistent_with_png(f.name)
+
+
+@pytest.mark.skipif(not Chart._svg_output_available(),
+                    reason='command-line tool vl2svg is not available')
+def test_savechart_svg():
+    chart = make_chart()
+
+    with tempfile.NamedTemporaryFile(suffix='.svg') as f:
+        chart.savechart(f.name)
+        assert consistent_with_svg(f.name)
+
+
+def test_savechart_html():
+    chart = make_chart()
+
+    with tempfile.NamedTemporaryFile(suffix='.html') as f:
+        chart.savechart(f.name)
+        content = f.read()
+        if hasattr(content, 'decode'):
+            content = content.decode('utf-8')
+        assert content.strip().startswith('<!DOCTYPE html>')
+
+
+def test_savechart_json():
+    chart = make_chart()
+
+    with tempfile.NamedTemporaryFile(suffix='.json') as f:
+        chart.savechart(f.name)
+        assert chart.to_dict() == json.load(open(f.name))
+
+
 def test_encode_update():
     # Test that encode updates rather than overwrites
     chart1 = Chart().encode(x='blah:Q').encode(y='blah:Q')
diff --git a/altair/utils/tests/test_node.py b/altair/utils/tests/test_node.py
index b54083e1..e62ed9f1 100644
--- a/altair/utils/tests/test_node.py
+++ b/altair/utils/tests/test_node.py
@@ -1,48 +1,33 @@
 import tempfile
-from xml.etree import ElementTree
 
 import pytest
 import pandas as pd
 
-from altair.utils.node import savechart, savechart_available, NodeExecError
+from altair.utils.node import savechart, vl_cmd_available, consistent_with_png, consistent_with_svg
 from altair import Chart
 
 
-def consistent_with_png(filename):
-    """Check that the file has a header consistent with PNG"""
-    with open(filename, 'rb') as f:
-        arr = f.read()
-    return arr.startswith(b'\x89PNG\r\n')
-
+def make_chart():
+    data = pd.DataFrame({'x': range(10),
+                         'y': range(10)})
+    return Chart(data).mark_point().encode(x='x', y='y')
 
-def consistent_with_svg(filename):
-    """Check that the file has a structure consistent with SVG"""
-    try:
-        e = ElementTree.parse(filename)
-    except ElementTree.ParseError:
-        return False
-    else:
-        keys = e.getroot().keys()
-        return set(keys) == {'width', 'version', 'class', 'height'}
 
+@pytest.mark.skipif(not vl_cmd_available('vl2png'),
+                    reason='command-line tool vl2png is not available')
+def test_savechart_png():
+    chart = make_chart()
 
-@pytest.mark.skipif(not savechart_available(),
-                    reason='node tools are not available')
-def test_savechart():
-    data = pd.DataFrame({'x': range(10),
-                         'y': range(10)})
-    chart = Chart(data).mark_point().encode(x='x', y='y')
-    
     with tempfile.NamedTemporaryFile(suffix='.png') as f:
-        try:
-            savechart(chart, f.name)
-        except NodeExecError:
-            pytest.skip('png failed due to improper nodejs setup')
+        savechart(chart, f.name)
         assert consistent_with_png(f.name)
 
+
+@pytest.mark.skipif(not vl_cmd_available('vl2svg'),
+                    reason='command-line tool vl2svg is not available')
+def test_savechart_svg():
+    chart = make_chart()
+
     with tempfile.NamedTemporaryFile(suffix='.svg') as f:
-        try:
-            savechart(chart, f.name)
-        except NodeExecError:
-            pytest.skip('svg failed due to improper nodejs setup')
+        savechart(chart, f.name)
         assert consistent_with_svg(f.name)