mirror of
https://github.com/xonsh/xonsh.git
synced 2025-03-04 08:24:40 +01:00
0806b0ff48
- also, removes reliance on the canvas object in favor of simple high-level figure methods
134 lines
4.1 KiB
Python
134 lines
4.1 KiB
Python
"""Matplotlib hooks, for what its worth."""
|
|
from io import BytesIO
|
|
import shutil
|
|
|
|
import numpy as np
|
|
import matplotlib
|
|
import matplotlib.pyplot as plt
|
|
|
|
from xonsh.tools import print_color, ON_WINDOWS
|
|
|
|
|
|
XONTRIB_MPL_MINIMAL_DEFAULT = True
|
|
|
|
|
|
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
|
|
|
|
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
|
|
"""
|
|
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
|
|
"""
|
|
# 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)
|
|
|
|
# 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
|
|
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
|
|
|
|
|
|
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)
|
|
|
|
|
|
def show():
|
|
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_tight_array(fig, w, h, minimal)
|
|
s = buf_to_color_str(buf)
|
|
print_color(s)
|