from __future__ import annotations import uuid from pyqtgraph import TextItem, mkPen, GraphicsObject, mkBrush from numpy import log10 from gui_qt.Qt import QtCore, QtGui from gui_qt.lib.pg_objects import LogInfiniteLine, PlotItem __all__ = ['LineObject', 'MultipointObject', 'RectangleObject', 'TextObject', 'EllipseObject'] class BaseObject: def __init__(self, color=None): self.id = str(uuid.uuid4()) self.color = color self.drawing = None def __str__(self): raise NotImplementedError class LineObject(BaseObject): def __init__(self, pos: float, angle: float, **kwargs): super().__init__(**kwargs) self.pos = pos self.angle = angle self.drawing = LogInfiniteLine(pos=self.pos, angle=self.angle, pen=mkPen(color=self.color.rgb())) def __str__(self): return f'{"x" if self.angle==90 else "y"}={self.pos}' class MultipointObject(BaseObject): def __init__(self, pts: list, closed: bool = False, **kwargs): super().__init__(**kwargs) x, y = zip(*pts) self.closed = closed if self.closed: x += (x[0],) y += (y[0],) self._x = x self._y = y self.drawing = PlotItem(x=self._x, y=self._y, pen=mkPen(color=self.color.rgb())) def __str__(self): return f'{len(self._y)-int(self.closed)}-pts' class Rectangle(GraphicsObject): # adapted from pyqtgraph example on custom objects def __init__(self, left: tuple, right: tuple): GraphicsObject.__init__(self) self.p0 = left self.delta = [right[0]-self.p0[0], right[1]-self.p0[1]] self.x_disp = self.p0[0] self.dx = self.delta[0] self.y_disp = self.p0[1] self.dy = self.delta[1] self.log_mode = [False, False] self.generatePicture() def generatePicture(self): ## pre-computing a QPicture object allows paint() to run much more quickly, ## rather than re-drawing the shapes every time. self.picture = QtGui.QPicture() p = QtGui.QPainter(self.picture) p.setPen(mkPen('w')) p.setBrush(mkBrush('r')) p.drawRect(QtCore.QRectF(self.x_disp, self.y_disp, self.dx, self.dy)) p.end() def paint(self, p, *_): p.drawPicture(0, 0, self.picture) def boundingRect(self): ## boundingRect _must_ indicate the entire area that will be drawn on ## or else we will get artifacts and possibly crashing. ## (in this case, QPicture does all the work of computing the bouning rect for us) return QtCore.QRectF(self.picture.boundingRect()) def setLogMode(self, x_mode, y_mode): if self.log_mode == [x_mode, y_mode]: return if self.log_mode[0] != x_mode: if x_mode: x0 = self.p0[0] xmax = log10(x0+self.delta[0]) self.x_disp = log10(x0) self.dx = xmax - self.x_disp else: self.x_disp = self.p0[0] self.dx = self.delta[0] if self.log_mode[1] != y_mode: if y_mode: y0 = self.p0[1] ymax = log10(y0 + self.delta[1]) self.y_disp = log10(y0) self.dy = ymax - self.y_disp else: self.y_disp = self.p0[1] self.dy = self.delta[1] self.generatePicture() self.log_mode = [x_mode, y_mode] self.informViewBoundsChanged() class RectangleObject(BaseObject): def __init__(self, left: tuple, right: tuple, **kwargs): super().__init__(**kwargs) self.left = left self.right = right self.drawing = Rectangle(left, right) def __str__(self): return f'Rectangle {self.left} to {self.right}' class TextObject(BaseObject): def __init__(self, text, pos, **kwargs): super().__init__(**kwargs) self.pos = pos self.drawing = TextItem(text, anchor=(0.5, 0.5), color=self.color.rgb()) self.drawing.setPos(*pos) def __str__(self): return f'{self.drawing.toPlainText()} at {self.pos}' class EllipseObject(BaseObject): def __init__(self, center: tuple, axes: tuple, **kwargs): super().__init__(**kwargs) self.center = center self.axes = axes self.drawing = Ellipse(center, axes) def __str__(self): return f'Ellipse at {self.center}' class Ellipse(GraphicsObject): # adapted from pyqtgraph example on custom objects def __init__(self, center: tuple, axes: tuple): GraphicsObject.__init__(self) self.center = center self.ax = axes self.center_disp = self.center self.ax_disp = self.ax self.log_mode = [False, False] self.generatePicture() def generatePicture(self): ## pre-computing a QPicture object allows paint() to run much more quickly, ## rather than re-drawing the shapes every time. self.picture = QtGui.QPicture() p = QtGui.QPainter(self.picture) p.setPen(mkPen('w')) p.setBrush(mkBrush('r')) p.drawEllipse(QtCore.QPointF(*self.center_disp), *self.ax_disp) p.end() def paint(self, p, *_): p.drawPicture(0, 0, self.picture) def boundingRect(self): ## boundingRect _must_ indicate the entire area that will be drawn on ## or else we will get artifacts and possibly crashing. ## (in this case, QPicture does all the work of computing the bouning rect for us) return QtCore.QRectF(self.picture.boundingRect()) def setLogMode(self, x_mode, y_mode): if self.log_mode == [x_mode, y_mode]: return if self.log_mode[0] != x_mode: if x_mode: x0 = self.p0[0] xmax = log10(x0+self.delta[0]) self.x_disp = log10(x0) self.dx = xmax - self.x_disp else: self.x_disp = self.p0[0] self.dx = self.delta[0] if self.log_mode[1] != y_mode: if y_mode: y0 = self.p0[1] ymax = log10(y0 + self.delta[1]) self.y_disp = log10(y0) self.dy = ymax - self.y_disp else: self.y_disp = self.p0[1] self.dy = self.delta[1] self.generatePicture() self.log_mode = [x_mode, y_mode] self.informViewBoundsChanged() def get_state(self): return {'center': self.center, 'axes': self.ax}