"""
Plot base classes.
"""
__all__ = [
'SceneNode',
'SimpleArtist',
'SceneGraph',
'ListView',
'LineBundle',
'plot_line',
]
[docs]class SceneNode:
"""An element of a figure."""
# member variables (will be overriden per instance if deviating):
name = None
parent = None
shown = False # if this element is currently drawn
_enabled = True # whether this element (+children) should be drawn
_figure = None # the matplotlib figure we should draw on
items = () # child nodes
lines = None # drawn matplotlib figure elements
# public API:
[docs] def enable(self, enabled=True):
"""Enable/disable the element individually."""
if self._enabled != enabled:
self._enabled = enabled
self.render()
self.draw_idle()
[docs] def enabled(self):
"""Check whether this element should be drawn."""
return self._enabled and (
self.parent is None or self.parent.enabled())
[docs] def node(self, name):
"""Find and return child node by name."""
for item in self.items:
if item.name == name:
return item
# Prevent exceptions in code like `graph.node(NAME).enable(True)`
# after destruction of the graph:
return NullNode()
[docs] def invalidate(self):
"""Mark drawn state as stale and redraw if needed."""
if self.enabled() and self.shown:
self._update()
@property
def figure(self):
return self._figure or self.parent.figure
# private, should be called via the scene tree only:
[docs] def render(self, show=True):
"""Draw or remove this node (and all children) from the figure."""
show = self.enabled() and show
shown = self.shown
if show and not shown:
self._draw()
elif not show and shown:
self._erase()
self.shown = show
[docs] def destroy(self):
"""
Cleanup all allocated ressources and disconnect signals. Called when
the SceneNode is removed from the graph and not needed anymore.
"""
self.on_clear_figure()
# overrides, private, must only be called from `render`:
def _draw(self):
"""Plot the element in a newly created scene."""
def _erase(self):
"""Remove the element from the scene."""
def _update(self):
"""Update existing plot."""
# Used internally, must override in root node!
[docs] def draw_idle(self):
"""Let the canvas know that it has to redraw."""
self.parent.draw_idle()
class NullNode(SceneNode):
"""Non-existent node."""
_enabled = False
def __bool__(self):
return False
def draw_idle(self):
pass
[docs]class SimpleArtist(SceneNode):
"""Delegates to draw function that returns a list of matplotlib artists."""
def __init__(self, name, artist, *args, **kwargs):
self.name = name
self.artist = artist
self.args = args
self.kwargs = kwargs
def _draw(self):
self.lines = LineBundle([
self.artist(ax, *self.args, **self.kwargs)
for ax in self.figure.axes
])
def _erase(self):
self.lines.remove()
def _update(self):
self.lines = self.lines.redraw()
[docs]class SceneGraph(SceneNode):
"""A scene element that is composed of multiple elements."""
def __init__(self, name, items=(), figure=None):
self.name = name
self.items = list(items)
self._figure = figure
self._adopt(items)
def _adopt(self, items):
for item in items:
item.parent = self
# overrides
def _draw(self):
for item in self.items:
item.render(True)
def _erase(self):
for item in self.items:
item.render(False)
def _update(self):
for item in self.items:
item.invalidate()
# manage items:
[docs] def add(self, *items):
"""Extend by several children."""
self.extend(items)
[docs] def extend(self, items):
"""Extend by several children."""
self.items.extend(items)
self._adopt(items)
if self.shown:
for item in items:
item.render(True)
self.draw_idle()
[docs] def insert(self, index, item):
self.items.insert(index, item)
self._adopt((item,))
if self.shown:
item.render(True)
self.draw_idle()
[docs] def pop(self, item):
"""Remove and hide one item (by value)."""
item.render(False)
item.destroy()
self.items.remove(item)
self.draw_idle()
[docs] def destroy(self):
for item in self.items:
item.destroy()
self.items = ()
self.shown = False
[docs]class ListView(SceneGraph):
def __init__(self, name, model, fn, *args, **kwargs):
super().__init__(name)
self.model = model
self.fn = fn
self.args = args
self.kwargs = kwargs
for idx, item in enumerate(model):
self._add(idx, item)
model.inserted.connect(self._add)
model.removed.connect(self._rm)
model.changed.connect(self._chg)
def _add(self, idx, item):
# This handles items with `name` or `elem` attributes and allows
# `item.elem` to be either a name string or `Element` object:
elem = getattr(item, 'elem', item)
name = getattr(elem, 'name', elem)
args = self.args + (item,)
node = SimpleArtist(str(name), self.fn, *args, **self.kwargs)
self.insert(idx, node)
def _rm(self, idx):
self.pop(self.items[idx])
def _chg(self, idx, val):
self._rm(idx)
self._add(idx, val)
[docs] def destroy(self):
self.model.inserted.disconnect(self._add)
self.model.removed.disconnect(self._rm)
self.model.changed.disconnect(self._chg)
super().destroy()
[docs]class LineBundle(list):
__slots__ = ()
[docs] def remove(self):
for line in self:
line.remove()
self.clear()
[docs] def redraw(self):
self[:] = [
line.redraw() if hasattr(line, 'redraw') else line
for line in self
]
return self
[docs]class plot_line:
"""Plot a single line using an fetch function."""
def __init__(self, ax, get_xydata, **style):
self._get_xydata = get_xydata
self.line, = ax.plot(*get_xydata(), **style)
[docs] def redraw(self):
self.line.set_data(*self._get_xydata())
return self
[docs] def remove(self):
self.line.remove()