Source code for schrodinger.ui.qt.figureoptions

"""
Module that provides a GUI-based editor for matplotlib's figure options

Copyright © 2009 Pierre Raybaut
Licensed under the terms of the MIT License
see the mpl licenses directory for a copy of the license

Additional changes copyright Schrodinger, LLC. All rights reserved.
"""

import os.path as osp

from matplotlib import markers
from matplotlib.backends.qt_compat import QtGui
from matplotlib.colors import to_hex

from schrodinger.Qt import IS_PYQT6
if not IS_PYQT6:
    # formlayout isn't compatible with Qt 6.  See PANEL-20903.
    import formlayout


def get_icon(name):
    import matplotlib
    basedir = osp.join(matplotlib.get_data_path(), 'images')
    return QtGui.QIcon(osp.join(basedir, name))


LINESTYLES = {
    '-': 'Solid',
    '--': 'Dashed',
    '-.': 'DashDot',
    ':': 'Dotted',
    'steps': 'Steps',
    'none': 'None',
}

MARKERS = markers.MarkerStyle.markers


def _get_scale_options(default, allow_log, allow_linear):
    scale_options = [default]
    if allow_linear:
        scale_options.append('linear')
    if allow_log:
        scale_options.append('log')
    return scale_options


def figure_edit(axes, parent=None, allow_log=True, allow_linear=True):
    """Edit matplotlib figure options"""
    sep = (None, None)  # separator

    has_curve = len(axes.get_lines()) > 0

    xscale_options = _get_scale_options(axes.get_xscale(), allow_log,
                                        allow_linear)
    yscale_options = _get_scale_options(axes.get_yscale(), allow_log,
                                        allow_linear)

    # Get / General
    xmin, xmax = axes.get_xlim()
    ymin, ymax = axes.get_ylim()
    general = [('Title', axes.get_title()), sep, (None, "<b>X-Axis</b>"),
               ('Min', xmin), ('Max', xmax), ('Label', axes.get_xlabel()),
               ('Scale', xscale_options), sep, (None, "<b>Y-Axis</b>"),
               ('Min', ymin), ('Max', ymax), ('Label', axes.get_ylabel()),
               ('Scale', yscale_options)]

    if has_curve:
        # Get / Curves
        linedict = {}
        for line in axes.get_lines():
            label = line.get_label()
            if label == '_nolegend_':
                continue
            linedict[label] = line
        curves = []
        linestyles = list(LINESTYLES.items())
        markers = list(MARKERS.items())
        curvelabels = sorted(list(linedict))
        for label in curvelabels:
            line = linedict[label]
            curvedata = [
                ('Label', label),
                sep,
                (None, '<b>Line</b>'),
                ('Style', [line.get_linestyle()] + linestyles),
                ('Width', line.get_linewidth()),
                ('Color', to_hex(line.get_color())),
                sep,
                (None, '<b>Marker</b>'),
                ('Style', [line.get_marker()] + markers),
                ('Size', line.get_markersize()),
                ('Facecolor', to_hex(line.get_markerfacecolor())),
                ('Edgecolor', to_hex(line.get_markeredgecolor())),
            ]
            curves.append([curvedata, label, ""])

    datalist = [(general, "Axes", "")]
    if has_curve:
        datalist.append((curves, "Curves", ""))

    def apply_callback(data, widgets=None):
        """This function will be called to apply changes"""
        if has_curve:
            general, curves = data
        else:
            general, = data

        # Set / General
        title, xmin, xmax, xlabel, xscale, ymin, ymax, ylabel, yscale = general

        if axes.get_xscale() != xscale:
            axes.set_xscale(xscale)
        if axes.get_yscale() != yscale:
            axes.set_yscale(yscale)

        axes.set_title(title)
        axes.set_xlim(xmin, xmax)
        axes.set_xlabel(xlabel)
        axes.set_ylim(ymin, ymax)
        axes.set_ylabel(ylabel)

        if has_curve:
            # Set / Curves
            for index, curve in enumerate(curves):
                line = linedict[curvelabels[index]]
                label, linestyle, linewidth, color, \
                    marker, markersize, markerfacecolor, markeredgecolor = curve
                line.set_label(label)
                line.set_linestyle(linestyle)
                line.set_linewidth(linewidth)
                line.set_color(color)
                if marker != 'none':
                    line.set_marker(marker)
                    line.set_markersize(markersize)
                    line.set_markerfacecolor(markerfacecolor)
                    line.set_markeredgecolor(markeredgecolor)

            # Recalculate the color and style of legend artists - note, this
            # does not work for multiple legends on the same figure.  But it's
            # an improvement over the current situation of not working for any
            # legends.
            figure = axes.get_figure()
            all_axes = figure.get_axes()
            for legend in figure.legends:
                # Need to gather all the legend items from all the subplots
                all_handles = []
                all_labels = []
                for one_axes in all_axes:
                    handles, labels = one_axes.get_legend_handles_labels()
                    all_handles.extend(handles)
                    all_labels.extend(labels)
                # Location and title are altered by the _init_legend_box
                # function, so preserve the current values
                loc = legend._loc
                title = legend.get_title().get_text()
                if title == 'None':
                    title = None
                # This function recalculates colors and styles
                legend._init_legend_box(all_handles, all_labels)
                legend._loc = loc
                legend.set_title(title)

            # Redraw
        figure = axes.get_figure()
        figure.canvas.draw()

    if not IS_PYQT6:
        data = formlayout.fedit(datalist,
                                title="Figure options",
                                parent=parent,
                                icon=get_icon('qt4_editor_options.svg'),
                                apply=apply_callback)
        if data is not None:
            apply_callback(data)