widget.py
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()