Source code for xontrib.mplhooks

"""Matplotlib hooks, for what its worth."""
from io import BytesIO
import shutil

import numpy as np
import matplotlib
import matplotlib.pyplot as plt

from import print_color, ON_WINDOWS

    # Use iterm2_tools as an indicator for the iterm2 terminal emulator
    from iterm2_tools.images import display_image_bytes
except ImportError:
    _use_iterm = False
    _use_iterm = True


def _get_buffer(fig, **kwargs):
    b = BytesIO()
    fig.savefig(b, **kwargs)
    return b

[docs]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)
[docs]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 = 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) # reduce 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 array = figure_to_rgb_array(fig, shape=(height, width, 4)) 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
[docs]def buf_to_color_str(buf): """Converts an RGB array to a xonsh color string.""" space = ' ' pix = '{{bg#{0:02x}{1:02x}{2:02x}}} ' pixels = [] for h in range(buf.shape[0]): last = None for w in range(buf.shape[1]): rgb = buf[h, w] if last is not None and (last == rgb).all(): pixels.append(space) else: pixels.append(pix.format(*rgb)) last = rgb pixels.append('{NO_COLOR}\n') pixels[-1] = pixels[-1].rstrip() return ''.join(pixels)
[docs]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()))
[docs]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() 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)