CommonMPL.py

#!/usr/bin/env python3
"""A TraitsUI editor that embeds a matplotlib figure in a Qt5 application.

Originally a Qt adaptation by Pierre Haessig (March 2014) of Gael Varoquaux's
tutorial on integrating matplotlib with TraitsUI, based on Qt code shared by
Didrik Pinte (May 2012):

* https://docs.enthought.com/traitsui/tutorials/traits_ui_scientific_app.html
* https://gist.github.com/pierre-haessig/9838326

Import order matters here. ``pyface.qt`` is imported first so that it resolves
and caches the Qt binding (PyQt5, PySide2, ...); matplotlib then reuses the
already-imported binding instead of independently picking a different one,
which would crash the moment the two tried to share a widget.
"""

# isort: off
# Order is deliberate -- see the module docstring. pyface.qt must resolve the
# Qt binding before matplotlib looks for one. Do not let a formatter sort these.
from pyface.qt import QtGui

from matplotlib.backends.backend_qtagg import (
    FigureCanvasQTAgg as FigureCanvas,
    NavigationToolbar2QT as NavigationToolbar,
)

# isort: on

from traitsui.api import Handler
from traitsui.basic_editor_factory import BasicEditorFactory
from traitsui.qt.editor import Editor

__all__ = ["MPLFigureEditor", "MPLInitHandler"]


class _MPLFigureEditor(Editor):
    """Wraps a :class:`~matplotlib.figure.Figure` trait in a Qt canvas."""

    scrollable = True

    def init(self, parent):
        self._canvas = None
        self.control = self._create_canvas(parent)
        self.set_tooltip()

    def update_editor(self):
        """Redraw the canvas when the underlying figure trait changes."""
        if self._canvas is not None:
            self._canvas.draw_idle()

    def _create_canvas(self, parent):
        """Build a frame holding the MPL canvas and its navigation toolbar."""
        frame = QtGui.QWidget()
        mpl_canvas = FigureCanvas(self.value)
        mpl_canvas.setParent(frame)
        self._canvas = mpl_canvas
        mpl_toolbar = NavigationToolbar(mpl_canvas, frame)

        vbox = QtGui.QVBoxLayout()
        vbox.addWidget(mpl_canvas)
        vbox.addWidget(mpl_toolbar)
        frame.setLayout(vbox)

        return frame


class MPLFigureEditor(BasicEditorFactory):
    """Editor factory: ``Item('figure', editor=MPLFigureEditor())``."""

    klass = _MPLFigureEditor


class MPLInitHandler(Handler):
    """Calls ``mpl_setup()`` on the view object once its controls exist.

    Attach this to a ``View`` when you need to register matplotlib event
    handlers (pickers, selectors) that require a live canvas.
    """

    def init(self, info):
        info.object.mpl_setup()
        return True