Merge pull request #1790 from derekbrokeit/xontrib-mpl_copy-fig

xontrib-mpl makes a copy of the current figure before modifying it

.circle_miniconda.sh

@@ -1,10 +1,10 @@
 if [[ $CIRCLE_NODE_INDEX == 0 ]]
 then
-    conda create -q -n test_env python=3.4 pygments prompt_toolkit ply pytest pytest-timeout psutil
+    conda create -q -n test_env python=3.4 pygments prompt_toolkit ply pytest pytest-timeout psutil numpy matplotlib
 fi
 
 
 if [[ $CIRCLE_NODE_INDEX == 1 ]]
 then
-    conda create -q -n test_env python=3.5 pygments prompt_toolkit ply pytest pytest-timeout psutil
+    conda create -q -n test_env python=3.5 pygments prompt_toolkit ply pytest pytest-timeout psutil numpy matplotlib
 fi

news/xontrib-mpl_copy-fig.rst

@@ -0,0 +1,16 @@
+**Added:** 
+
+* $XONTRIB_MPL_MINIMAL environment variable can be set to change if plots are minimalist or as-seen
+* xontrib-mpl now supports iTerm2 inline image display if iterm2_tools python package is installed
+
+**Changed:** None
+
+**Deprecated:** None
+
+**Removed:** None
+
+**Fixed:** 
+
+* xontrib-mpl now preserves the figure and does not permanently alter it for viewing
+
+**Security:** None

requirements-tests.txt

@@ -10,3 +10,5 @@ pygments==2.1.3
 codecov==2.0.5
 flake8==2.6.2
 coverage==4.2
+numpy==1.10.4
+matplotlib==1.5.1

tests/test_mpl.py

@@ -0,0 +1,104 @@
+import pytest
+
+# make sure to skip these tests entirely if numpy/matplotlib are not present
+np = pytest.importorskip("numpy")
+matplotlib = pytest.importorskip("matplotlib")
+plt = pytest.importorskip("matplotlib.pyplot")
+
+from xontrib import mplhooks
+
+
+# some default settings that are temporarily changed by mpl
+FONT_SIZE    = 22
+FACE_COLOR   = (0.0, 1.0, 0.0, 1.0)
+DPI          = 80
+
+
+def create_figure():
+    """Simply create a figure with the default settings"""
+    f, ax = plt.subplots()
+    ax.plot(np.arange(20), np.arange(20))
+    # set the figure parameters such that mpl will require changes
+    f.set_facecolor(FACE_COLOR)
+    f.dpi = DPI
+    matplotlib.rcParams.update({'font.size': FONT_SIZE})
+    return f
+
+
+def test_mpl_preserve_font_size():
+    """Make sure that matplotlib preserves font size settings"""
+    f = create_figure()
+    width, height = f.canvas.get_width_height()
+    s = mplhooks.figure_to_tight_array(f, 0.5*width, 0.5*height, True)
+    exp = FONT_SIZE
+    obs = matplotlib.rcParams['font.size']
+    plt.close(f)
+    assert exp == obs
+
+
+def test_mpl_preserve_face_color():
+    """Make sure that the figure preserves face color settings"""
+    f = create_figure()
+    width, height = f.canvas.get_width_height()
+    s = mplhooks.figure_to_tight_array(f, 0.5*width, 0.5*height, True)
+    exp = FACE_COLOR
+    obs = f.get_facecolor()
+    plt.close(f)
+    assert exp == obs
+
+
+def test_mpl_preserve_width():
+    """Make sure that the figure preserves width settings"""
+    f = create_figure()
+    width, height = f.canvas.get_width_height()
+    s = mplhooks.figure_to_tight_array(f, 0.5*width, 0.5*height, True)
+    exp = width
+    newwidth, newheight = f.canvas.get_width_height()
+    obs = newwidth
+    plt.close(f)
+    assert exp == obs
+
+
+def test_mpl_preserve_height():
+    """Make sure that the figure preserves height settings"""
+    f = create_figure()
+    width, height = f.canvas.get_width_height()
+    s = mplhooks.figure_to_tight_array(f, 0.5*width, 0.5*height, True)
+    exp = height
+    newwidth, newheight = f.canvas.get_width_height()
+    obs = newheight
+    plt.close(f)
+    assert exp == obs
+
+
+def test_mpl_preserve_dpi():
+    """Make sure that the figure preserves height settings"""
+    f = create_figure()
+    width, height = f.canvas.get_width_height()
+    s = mplhooks.figure_to_tight_array(f, 0.5*width, 0.5*height, False)
+    exp = DPI
+    obs = f.dpi
+    plt.close(f)
+    assert exp == obs
+
+
+def test_mpl_preserve_image_tight():
+    """Make sure that the figure preserves height settings"""
+    f = create_figure()
+    exp = mplhooks.figure_to_rgb_array(f)
+    width, height = f.canvas.get_width_height()
+    s = mplhooks.figure_to_tight_array(f, 0.5*width, 0.5*height, True)
+    obs = mplhooks.figure_to_rgb_array(f)
+    plt.close(f)
+    assert np.all(exp == obs)
+
+
+def test_mpl_preserve_standard():
+    """Make sure that the figure preserves height settings"""
+    f = create_figure()
+    exp = mplhooks.figure_to_rgb_array(f)
+    width, height = f.canvas.get_width_height()
+    s = mplhooks.figure_to_tight_array(f, 0.5*width, 0.5*height, False)
+    obs = mplhooks.figure_to_rgb_array(f)
+    plt.close(f)
+    assert np.all(exp == obs)

xontrib/mplhooks.py

@@ -1,4 +1,5 @@
 """Matplotlib hooks, for what its worth."""
+from io import BytesIO
 import shutil
 
 import numpy as np
@@ -7,33 +8,114 @@ import matplotlib.pyplot as plt
 
 from xonsh.tools import print_color, ON_WINDOWS
 
+try:
+    # Use iterm2_tools as an indicator for the iterm2 terminal emulator
+    from iterm2_tools.images import display_image_bytes
+except ImportError:
+    _use_iterm = False
+else:
+    _use_iterm = True
 
-def figure_to_rgb_array(fig, width, height):
-    """Converts figure to a numpy array of rgb values
+XONTRIB_MPL_MINIMAL_DEFAULT = True
 
-    Forked from http://www.icare.univ-lille1.fr/wiki/index.php/How_to_convert_a_matplotlib_figure_to_a_numpy_array_or_a_PIL_image
+
+def _get_buffer(fig, **kwargs):
+    b = BytesIO()
+    fig.savefig(b, **kwargs)
+    b.seek(0)
+    return b
+
+
+def figure_to_rgb_array(fig, shape=None):
+    """Converts figure to a numpy array
+
+    Parameters
+    ----------
+    fig : matplotlib.figure.Figure
+        the figure to be plotted
+    shape : iterable
+        with the shape of the output array. by default this attempts to use the
+        pixel height and width of the figure
+
+
+    Returns
+    -------
+    array : np.ndarray
+        An RGBA array of the image represented by the figure.
+
+    Note: the method will throw an exception if the given shape is wrong.
     """
+    array = np.frombuffer(_get_buffer(fig, dpi=fig.dpi, format='raw').read(), dtype='uint8')
+    if shape is None:
+        w, h = fig.canvas.get_width_height()
+        shape = (h, w, 4)
+    return array.reshape(*shape)
+
+
+def figure_to_tight_array(fig, width, height, minimal=True):
+    """Converts figure to a numpy array of rgb values of tight value
+
+    Parameters
+    ----------
+    fig : matplotlib.figure.Figure
+        the figure to be plotted
+    width : int
+        pixel width of the final array
+    height : int
+        pixel height of the final array
+    minimal : bool
+        whether or not to reduce the output array to minimized margins/whitespace
+        text is also eliminated
+
+    Returns
+    -------
+    array : np.ndarray
+        An RGBA array of the image represented by the figure.
+    """
+    # store the properties of the figure in order to restore it
     w, h = fig.canvas.get_width_height()
-    dpi = fig.get_dpi()
-    fig.set_size_inches(width/dpi, height/dpi, forward=True)
-    width, height = fig.canvas.get_width_height()
-    ax = fig.gca()
-    ax.set_xticklabels([])
-    ax.set_yticklabels([])
-    fig.set_tight_layout(True)
-    fig.set_frameon(False)
-    fig.set_facecolor('w')
-    font_size = matplotlib.rcParams['font.size']
-    matplotlib.rcParams.update({'font.size': 1})
+    dpi_fig = fig.dpi
+    if minimal:
+        # perform reversible operations to produce an optimally tight layout
+        dpi = dpi_fig
+        subplotpars = {
+                k: getattr(fig.subplotpars, k)
+                for k in ['wspace', 'hspace', 'bottom', 'top', 'left', 'right']
+                }
+
+        # set the figure dimensions to the terminal size
+        fig.set_size_inches(width/dpi, height/dpi, forward=True)
+        width, height = fig.canvas.get_width_height()
+
+        # remove all space between subplots
+        fig.subplots_adjust(wspace=0, hspace=0)
+        # move all subplots to take the entirety of space in the figure
+        # leave only one line for top and bottom
+        fig.subplots_adjust(bottom=1/height, top=1-1/height, left=0, right=1)
+
+        # redeuce font size in order to reduce text impact on the image
+        font_size = matplotlib.rcParams['font.size']
+        matplotlib.rcParams.update({'font.size': 0})
+    else:
+        dpi = min([width * fig.dpi // w, height * fig.dpi // h])
+        fig.dpi = dpi
+        width, height = fig.canvas.get_width_height()
 
     # Draw the renderer and get the RGB buffer from the figure
-    fig.canvas.draw()
-    buf = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8)
-    buf.shape = (height, width, 3)
+    array = figure_to_rgb_array(fig, shape=(height, width, 4))
 
-    # clean up and return
-    matplotlib.rcParams.update({'font.size': font_size})
-    return buf
+    if minimal:
+        # cleanup after tight layout
+        # clean up rcParams
+        matplotlib.rcParams.update({'font.size': font_size})
+
+        # reset the axis positions and figure dimensions
+        fig.set_size_inches(w/dpi, h/dpi, forward=True)
+        fig.subplots_adjust(**subplotpars)
+    else:
+        fig.dpi = dpi_fig
+
+    return array
 
 
 def buf_to_color_str(buf):
@@ -55,12 +137,32 @@ def buf_to_color_str(buf):
     return ''.join(pixels)
 
 
+def display_figure_with_iterm2(fig):
+    """Displays a matplotlib figure using iterm2 inline-image escape sequence.
+
+    Parameters
+    ----------
+    fig : matplotlib.figure.Figure
+        the figure to be plotted
+    """
+    print(display_image_bytes(_get_buffer(fig, format='png', dpi=fig.dpi).read()))
+
+
 def show():
+    '''Run the mpl display sequence by printing the most recent figure to console'''
+    try:
+        minimal = __xonsh_env__['XONTRIB_MPL_MINIMAL']
+    except KeyError:
+        minimal = XONTRIB_MPL_MINIMAL_DEFAULT
     fig = plt.gcf()
-    w, h = shutil.get_terminal_size()
-    if ON_WINDOWS:
-        w -= 1  # @melund reports that win terminals are too thin
-    h -= 1  # leave space for next prompt
-    buf = figure_to_rgb_array(fig, w, h)
-    s = buf_to_color_str(buf)
-    print_color(s)
+    if _use_iterm:
+        display_figure_with_iterm2(fig)
+    else:
+        # Display the image using terminal characters to fit into the console
+        w, h = shutil.get_terminal_size()
+        if ON_WINDOWS:
+            w -= 1  # @melund reports that win terminals are too thin
+        h -= 1  # leave space for next prompt
+        buf = figure_to_tight_array(fig, w, h, minimal)
+        s = buf_to_color_str(buf)
+        print_color(s)