from datetime import datetime import pyqtgraph as pg # type: ignore from vnpy.trader.ui import QtGui, QtWidgets, QtCore from vnpy.trader.object import BarData from .manager import BarManager from .base import ( GREY_COLOR, WHITE_COLOR, CURSOR_COLOR, BLACK_COLOR, to_int, NORMAL_FONT ) from .axis import DatetimeAxis from .item import ChartItem pg.setConfigOptions(antialias=True) class ChartWidget(pg.PlotWidget): """""" MIN_BAR_COUNT = 100 def __init__(self, parent: QtWidgets.QWidget | None = None) -> None: """""" super().__init__(parent) self._manager: BarManager = BarManager() self._plots: dict[str, pg.PlotItem] = {} self._items: dict[str, ChartItem] = {} self._item_plot_map: dict[ChartItem, pg.PlotItem] = {} self._first_plot: pg.PlotItem | None = None self._cursor: ChartCursor | None = None self._right_ix: int = 0 # Index of most right data self._bar_count: int = self.MIN_BAR_COUNT # Total bar visible in chart self._init_ui() def _init_ui(self) -> None: """""" self.setWindowTitle("ChartWidget of VeighNa") self._layout: pg.GraphicsLayout = pg.GraphicsLayout() self._layout.setContentsMargins(10, 10, 10, 10) self._layout.setSpacing(0) self._layout.setBorder(color=GREY_COLOR, width=0.8) self._layout.setZValue(0) self.setCentralItem(self._layout) def _get_new_x_axis(self) -> DatetimeAxis: return DatetimeAxis(self._manager, orientation="bottom") def add_cursor(self) -> None: """""" if not self._cursor: self._cursor = ChartCursor( self, self._manager, self._plots, self._item_plot_map) def add_plot( self, plot_name: str, minimum_height: int = 80, maximum_height: int | None = None, hide_x_axis: bool = False ) -> None: """ Add plot area. """ # Create plot object plot: pg.PlotItem = pg.PlotItem(axisItems={"bottom": self._get_new_x_axis()}) plot.setMenuEnabled(False) plot.setClipToView(True) plot.hideAxis("left") plot.showAxis("right") plot.setDownsampling(mode="peak") plot.setRange(xRange=(0, 1), yRange=(0, 1)) plot.hideButtons() plot.setMinimumHeight(minimum_height) if maximum_height: plot.setMaximumHeight(maximum_height) if hide_x_axis: plot.hideAxis("bottom") if not self._first_plot: self._first_plot = plot # Connect view change signal to update y range function view: pg.ViewBox = plot.getViewBox() view.sigXRangeChanged.connect(self._update_y_range) view.setMouseEnabled(x=True, y=False) # Set right axis right_axis: pg.AxisItem = plot.getAxis("right") right_axis.setWidth(60) right_axis.tickFont = NORMAL_FONT # Connect x-axis link if self._plots: first_plot: pg.PlotItem = list(self._plots.values())[0] plot.setXLink(first_plot) # Store plot object in dict self._plots[plot_name] = plot # Add plot onto the layout self._layout.nextRow() self._layout.addItem(plot) def add_item( self, item_class: type[ChartItem], item_name: str, plot_name: str ) -> None: """ Add chart item. """ item: ChartItem = item_class(self._manager) self._items[item_name] = item plot: pg.PlotItem = self._plots.get(plot_name) plot.addItem(item) self._item_plot_map[item] = plot def get_plot(self, plot_name: str) -> pg.PlotItem: """ Get specific plot with its name. """ return self._plots.get(plot_name, None) def get_all_plots(self) -> list[pg.PlotItem]: """ Get all plot objects. """ return list(self._plots.values()) def clear_all(self) -> None: """ Clear all data. """ self._manager.clear_all() for item in self._items.values(): item.clear_all() if self._cursor: self._cursor.clear_all() def update_history(self, history: list[BarData]) -> None: """ Update a list of bar data. """ self._manager.update_history(history) for item in self._items.values(): item.update_history(history) self._update_plot_limits() self.move_to_right() def update_bar(self, bar: BarData) -> None: """ Update single bar data. """ self._manager.update_bar(bar) for item in self._items.values(): item.update_bar(bar) self._update_plot_limits() if self._right_ix >= (self._manager.get_count() - self._bar_count / 2): self.move_to_right() def _update_plot_limits(self) -> None: """ Update the limit of plots. """ for item, plot in self._item_plot_map.items(): min_value, max_value = item.get_y_range() plot.setLimits( xMin=-1, xMax=self._manager.get_count(), yMin=min_value, yMax=max_value ) def _update_x_range(self) -> None: """ Update the x-axis range of plots. """ max_ix: int = self._right_ix min_ix: int = self._right_ix - self._bar_count for plot in self._plots.values(): plot.setRange(xRange=(min_ix, max_ix), padding=0) def _update_y_range(self) -> None: """ Update the y-axis range of plots. """ if not self._first_plot: return view: pg.ViewBox = self._first_plot.getViewBox() view_range: list = view.viewRange() min_ix: int = max(0, int(view_range[0][0])) max_ix: int = min(self._manager.get_count(), int(view_range[0][1])) # Update limit for y-axis for item, plot in self._item_plot_map.items(): y_range: tuple = item.get_y_range(min_ix, max_ix) plot.setRange(yRange=y_range) def paintEvent(self, event: QtGui.QPaintEvent) -> None: """ Reimplement this method of parent to update current max_ix value. """ if not self._first_plot: return view: pg.ViewBox = self._first_plot.getViewBox() view_range: list = view.viewRange() self._right_ix = max(0, view_range[0][1]) super().paintEvent(event) def keyPressEvent(self, event: QtGui.QKeyEvent) -> None: """ Reimplement this method of parent to move chart horizontally and zoom in/out. """ Key = QtCore.Qt.Key if event.key() == Key.Key_Left: self._on_key_left() elif event.key() == Key.Key_Right: self._on_key_right() elif event.key() == Key.Key_Up: self._on_key_up() elif event.key() == Key.Key_Down: self._on_key_down() def wheelEvent(self, event: QtGui.QWheelEvent) -> None: """ Reimplement this method of parent to zoom in/out. """ delta: QtCore.QPoint = event.angleDelta() if delta.y() > 0: self._on_key_up() elif delta.y() < 0: self._on_key_down() def _on_key_left(self) -> None: """ Move chart to left. """ self._right_ix -= 1 self._right_ix = max(self._right_ix, self._bar_count) self._update_x_range() if self._cursor: self._cursor.move_left() self._cursor.update_info() def _on_key_right(self) -> None: """ Move chart to right. """ self._right_ix += 1 self._right_ix = min(self._right_ix, self._manager.get_count()) self._update_x_range() if self._cursor: self._cursor.move_right() self._cursor.update_info() def _on_key_down(self) -> None: """ Zoom out the chart. """ self._bar_count = int(self._bar_count * 1.2) self._bar_count = min(int(self._bar_count), self._manager.get_count()) self._update_x_range() if self._cursor: self._cursor.update_info() def _on_key_up(self) -> None: """ Zoom in the chart. """ self._bar_count = int(self._bar_count / 1.2) self._bar_count = max(int(self._bar_count), self.MIN_BAR_COUNT) self._update_x_range() if self._cursor: self._cursor.update_info() def move_to_right(self) -> None: """ Move chart to the most right. """ self._right_ix = self._manager.get_count() self._update_x_range() if self._cursor: self._cursor.update_info() class ChartCursor(QtCore.QObject): """""" def __init__( self, widget: ChartWidget, manager: BarManager, plots: dict[str, pg.GraphicsObject], item_plot_map: dict[ChartItem, pg.GraphicsObject] ) -> None: """""" super().__init__() self._widget: ChartWidget = widget self._manager: BarManager = manager self._plots: dict[str, pg.GraphicsObject] = plots self._item_plot_map: dict[ChartItem, pg.GraphicsObject] = item_plot_map self._x: int = 0 self._y: float = 0 self._plot_name: str = "" self._init_ui() self._connect_signal() def _init_ui(self) -> None: """""" self._init_line() self._init_label() self._init_info() def _init_line(self) -> None: """ Create line objects. """ self._v_lines: dict[str, pg.InfiniteLine] = {} self._h_lines: dict[str, pg.InfiniteLine] = {} self._views: dict[str, pg.ViewBox] = {} pen: QtGui.QPen = pg.mkPen(WHITE_COLOR) for plot_name, plot in self._plots.items(): v_line: pg.InfiniteLine = pg.InfiniteLine(angle=90, movable=False, pen=pen) h_line: pg.InfiniteLine = pg.InfiniteLine(angle=0, movable=False, pen=pen) view: pg.ViewBox = plot.getViewBox() for line in [v_line, h_line]: line.setZValue(0) line.hide() view.addItem(line) self._v_lines[plot_name] = v_line self._h_lines[plot_name] = h_line self._views[plot_name] = view def _init_label(self) -> None: """ Create label objects on axis. """ self._y_labels: dict[str, pg.TextItem] = {} for plot_name, plot in self._plots.items(): label: pg.TextItem = pg.TextItem( plot_name, fill=CURSOR_COLOR, color=BLACK_COLOR) label.hide() label.setZValue(2) label.setFont(NORMAL_FONT) plot.addItem(label, ignoreBounds=True) self._y_labels[plot_name] = label self._x_label: pg.TextItem = pg.TextItem( "datetime", fill=CURSOR_COLOR, color=BLACK_COLOR) self._x_label.hide() self._x_label.setZValue(2) self._x_label.setFont(NORMAL_FONT) plot.addItem(self._x_label, ignoreBounds=True) def _init_info(self) -> None: """ """ self._infos: dict[str, pg.TextItem] = {} for plot_name, plot in self._plots.items(): info: pg.TextItem = pg.TextItem( "info", color=CURSOR_COLOR, border=CURSOR_COLOR, fill=BLACK_COLOR ) info.hide() info.setZValue(2) info.setFont(NORMAL_FONT) plot.addItem(info) # , ignoreBounds=True) self._infos[plot_name] = info def _connect_signal(self) -> None: """ Connect mouse move signal to update function. """ self._widget.scene().sigMouseMoved.connect(self._mouse_moved) def _mouse_moved(self, evt: tuple) -> None: """ Callback function when mouse is moved. """ if not self._manager.get_count(): return # First get current mouse point pos: tuple = evt for plot_name, view in self._views.items(): rect = view.sceneBoundingRect() if rect.contains(pos): mouse_point = view.mapSceneToView(pos) self._x = to_int(mouse_point.x()) self._y = mouse_point.y() self._plot_name = plot_name break # Then update cursor component self._update_line() self._update_label() self.update_info() def _update_line(self) -> None: """""" for v_line in self._v_lines.values(): v_line.setPos(self._x) v_line.show() for plot_name, h_line in self._h_lines.items(): if plot_name == self._plot_name: h_line.setPos(self._y) h_line.show() else: h_line.hide() def _update_label(self) -> None: """""" bottom_plot: pg.PlotItem = list(self._plots.values())[-1] axis_width = bottom_plot.getAxis("right").width() axis_height = bottom_plot.getAxis("bottom").height() axis_offset: QtCore.QPointF = QtCore.QPointF(axis_width, axis_height) bottom_view: pg.ViewBox = list(self._views.values())[-1] bottom_right = bottom_view.mapSceneToView( bottom_view.sceneBoundingRect().bottomRight() - axis_offset ) for plot_name, label in self._y_labels.items(): if plot_name == self._plot_name: label.setText(str(self._y)) label.show() label.setPos(bottom_right.x(), self._y) else: label.hide() dt: datetime | None = self._manager.get_datetime(self._x) if dt: self._x_label.setText(dt.strftime("%Y-%m-%d %H:%M:%S")) self._x_label.show() self._x_label.setPos(self._x, bottom_right.y()) self._x_label.setAnchor((0, 0)) def update_info(self) -> None: """""" buf: dict = {} for item, plot in self._item_plot_map.items(): item_info_text: str = item.get_info_text(self._x) if plot not in buf: buf[plot] = item_info_text else: if item_info_text: buf[plot] += ("\n\n" + item_info_text) for plot_name, plot in self._plots.items(): plot_info_text: str = buf[plot] info: pg.TextItem = self._infos[plot_name] info.setText(plot_info_text) info.show() view: pg.ViewBox = self._views[plot_name] top_left = view.mapSceneToView(view.sceneBoundingRect().topLeft()) info.setPos(top_left) def move_right(self) -> None: """ Move cursor index to right by 1. """ if self._x == self._manager.get_count() - 1: return self._x += 1 self._update_after_move() def move_left(self) -> None: """ Move cursor index to left by 1. """ if self._x == 0: return self._x -= 1 self._update_after_move() def _update_after_move(self) -> None: """ Update cursor after moved by left/right. """ bar: BarData | None = self._manager.get_bar(self._x) if bar is None: return self._y = bar.close_price self._update_line() self._update_label() def clear_all(self) -> None: """ Clear all data. """ self._x = 0 self._y = 0 self._plot_name = "" for line in list(self._v_lines.values()) + list(self._h_lines.values()): line.hide() for label in list(self._y_labels.values()) + [self._x_label]: label.hide()