334 lines
9.3 KiB
Python
334 lines
9.3 KiB
Python
from abc import abstractmethod
|
|
|
|
import pyqtgraph as pg # type: ignore
|
|
|
|
from vnpy.trader.ui import QtCore, QtGui, QtWidgets
|
|
from vnpy.trader.object import BarData
|
|
|
|
from .base import BLACK_COLOR, UP_COLOR, DOWN_COLOR, PEN_WIDTH, BAR_WIDTH
|
|
from .manager import BarManager
|
|
|
|
|
|
class ChartItem(pg.GraphicsObject):
|
|
""""""
|
|
|
|
def __init__(self, manager: BarManager) -> None:
|
|
""""""
|
|
super().__init__()
|
|
|
|
self._manager: BarManager = manager
|
|
|
|
self._bar_picutures: dict[int, QtGui.QPicture | None] = {}
|
|
self._item_picuture: QtGui.QPicture | None = None
|
|
|
|
self._black_brush: QtGui.QBrush = pg.mkBrush(color=BLACK_COLOR)
|
|
|
|
self._up_pen: QtGui.QPen = pg.mkPen(
|
|
color=UP_COLOR, width=PEN_WIDTH
|
|
)
|
|
self._up_brush: QtGui.QBrush = pg.mkBrush(color=UP_COLOR)
|
|
|
|
self._down_pen: QtGui.QPen = pg.mkPen(
|
|
color=DOWN_COLOR, width=PEN_WIDTH
|
|
)
|
|
self._down_brush: QtGui.QBrush = pg.mkBrush(color=DOWN_COLOR)
|
|
|
|
self._rect_area: tuple[float, float] | None = None
|
|
|
|
# Very important! Only redraw the visible part and improve speed a lot.
|
|
self.setFlag(self.GraphicsItemFlag.ItemUsesExtendedStyleOption)
|
|
|
|
# Force update during the next paint
|
|
self._to_update: bool = False
|
|
|
|
@abstractmethod
|
|
def _draw_bar_picture(self, ix: int, bar: BarData) -> QtGui.QPicture:
|
|
"""
|
|
Draw picture for specific bar.
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def boundingRect(self) -> QtCore.QRectF:
|
|
"""
|
|
Get bounding rectangles for item.
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_y_range(self, min_ix: int | None = None, max_ix: int | None = None) -> tuple[float, float]:
|
|
"""
|
|
Get range of y-axis with given x-axis range.
|
|
|
|
If min_ix and max_ix not specified, then return range with whole data set.
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_info_text(self, ix: int) -> str:
|
|
"""
|
|
Get information text to show by cursor.
|
|
"""
|
|
pass
|
|
|
|
def update_history(self, history: list[BarData]) -> None:
|
|
"""
|
|
Update a list of bar data.
|
|
"""
|
|
self._bar_picutures.clear()
|
|
|
|
bars: list[BarData] = self._manager.get_all_bars()
|
|
|
|
for ix, _ in enumerate(bars):
|
|
self._bar_picutures[ix] = None
|
|
|
|
self.update()
|
|
|
|
def update_bar(self, bar: BarData) -> None:
|
|
"""
|
|
Update single bar data.
|
|
"""
|
|
ix: int | None = self._manager.get_index(bar.datetime)
|
|
if ix is None:
|
|
return
|
|
|
|
self._bar_picutures[ix] = None
|
|
|
|
self.update()
|
|
|
|
def update(self) -> None:
|
|
"""
|
|
Refresh the item.
|
|
"""
|
|
if self.scene():
|
|
self._to_update = True
|
|
self.scene().update()
|
|
|
|
def paint(
|
|
self,
|
|
painter: QtGui.QPainter,
|
|
opt: QtWidgets.QStyleOptionGraphicsItem,
|
|
w: QtWidgets.QWidget
|
|
) -> None:
|
|
"""
|
|
Reimplement the paint method of parent class.
|
|
|
|
This function is called by external QGraphicsView.
|
|
"""
|
|
rect: QtCore.QRectF = opt.exposedRect # type: ignore
|
|
|
|
min_ix: int = int(rect.left())
|
|
max_ix: int = int(rect.right())
|
|
max_ix = min(max_ix, len(self._bar_picutures))
|
|
|
|
rect_area: tuple = (min_ix, max_ix)
|
|
if (
|
|
self._to_update
|
|
or rect_area != self._rect_area
|
|
or not self._item_picuture
|
|
):
|
|
self._to_update = False
|
|
self._rect_area = rect_area
|
|
self._draw_item_picture(min_ix, max_ix)
|
|
|
|
if self._item_picuture:
|
|
self._item_picuture.play(painter)
|
|
|
|
def _draw_item_picture(self, min_ix: int, max_ix: int) -> None:
|
|
"""
|
|
Draw the picture of item in specific range.
|
|
"""
|
|
self._item_picuture = QtGui.QPicture()
|
|
painter: QtGui.QPainter = QtGui.QPainter(self._item_picuture)
|
|
|
|
for ix in range(min_ix, max_ix):
|
|
bar_picture: QtGui.QPicture | None = self._bar_picutures[ix]
|
|
|
|
if bar_picture is None:
|
|
bar: BarData | None = self._manager.get_bar(ix)
|
|
if bar is None:
|
|
continue
|
|
|
|
bar_picture = self._draw_bar_picture(ix, bar)
|
|
self._bar_picutures[ix] = bar_picture
|
|
|
|
bar_picture.play(painter)
|
|
|
|
painter.end()
|
|
|
|
def clear_all(self) -> None:
|
|
"""
|
|
Clear all data in the item.
|
|
"""
|
|
self._item_picuture = None
|
|
self._bar_picutures.clear()
|
|
self.update()
|
|
|
|
|
|
class CandleItem(ChartItem):
|
|
""""""
|
|
|
|
def __init__(self, manager: BarManager) -> None:
|
|
""""""
|
|
super().__init__(manager)
|
|
|
|
def _draw_bar_picture(self, ix: int, bar: BarData) -> QtGui.QPicture:
|
|
""""""
|
|
# Create objects
|
|
candle_picture: QtGui.QPicture = QtGui.QPicture()
|
|
painter: QtGui.QPainter = QtGui.QPainter(candle_picture)
|
|
|
|
# Set painter color
|
|
if bar.close_price >= bar.open_price:
|
|
painter.setPen(self._up_pen)
|
|
painter.setBrush(self._black_brush)
|
|
else:
|
|
painter.setPen(self._down_pen)
|
|
painter.setBrush(self._down_brush)
|
|
|
|
# Draw candle shadow
|
|
if bar.high_price > bar.low_price:
|
|
painter.drawLine(
|
|
QtCore.QPointF(ix, bar.high_price),
|
|
QtCore.QPointF(ix, bar.low_price)
|
|
)
|
|
|
|
# Draw candle body
|
|
if bar.open_price == bar.close_price:
|
|
painter.drawLine(
|
|
QtCore.QPointF(ix - BAR_WIDTH, bar.open_price),
|
|
QtCore.QPointF(ix + BAR_WIDTH, bar.open_price),
|
|
)
|
|
else:
|
|
rect: QtCore.QRectF = QtCore.QRectF(
|
|
ix - BAR_WIDTH,
|
|
bar.open_price,
|
|
BAR_WIDTH * 2,
|
|
bar.close_price - bar.open_price
|
|
)
|
|
painter.drawRect(rect)
|
|
|
|
# Finish
|
|
painter.end()
|
|
return candle_picture
|
|
|
|
def boundingRect(self) -> QtCore.QRectF:
|
|
""""""
|
|
min_price, max_price = self._manager.get_price_range()
|
|
rect: QtCore.QRectF = QtCore.QRectF(
|
|
0,
|
|
min_price,
|
|
len(self._bar_picutures),
|
|
max_price - min_price
|
|
)
|
|
return rect
|
|
|
|
def get_y_range(self, min_ix: int | None = None, max_ix: int | None = None) -> tuple[float, float]:
|
|
"""
|
|
Get range of y-axis with given x-axis range.
|
|
|
|
If min_ix and max_ix not specified, then return range with whole data set.
|
|
"""
|
|
min_price, max_price = self._manager.get_price_range(min_ix, max_ix)
|
|
return min_price, max_price
|
|
|
|
def get_info_text(self, ix: int) -> str:
|
|
"""
|
|
Get information text to show by cursor.
|
|
"""
|
|
bar: BarData | None = self._manager.get_bar(ix)
|
|
|
|
if bar:
|
|
words: list = [
|
|
"Date",
|
|
bar.datetime.strftime("%Y-%m-%d"),
|
|
"",
|
|
"Time",
|
|
bar.datetime.strftime("%H:%M"),
|
|
"",
|
|
"Open",
|
|
str(bar.open_price),
|
|
"",
|
|
"High",
|
|
str(bar.high_price),
|
|
"",
|
|
"Low",
|
|
str(bar.low_price),
|
|
"",
|
|
"Close",
|
|
str(bar.close_price)
|
|
]
|
|
text: str = "\n".join(words)
|
|
else:
|
|
text = ""
|
|
|
|
return text
|
|
|
|
|
|
class VolumeItem(ChartItem):
|
|
""""""
|
|
|
|
def __init__(self, manager: BarManager) -> None:
|
|
""""""
|
|
super().__init__(manager)
|
|
|
|
def _draw_bar_picture(self, ix: int, bar: BarData) -> QtGui.QPicture:
|
|
""""""
|
|
# Create objects
|
|
volume_picture: QtGui.QPicture = QtGui.QPicture()
|
|
painter: QtGui.QPainter = QtGui.QPainter(volume_picture)
|
|
|
|
# Set painter color
|
|
if bar.close_price >= bar.open_price:
|
|
painter.setPen(self._up_pen)
|
|
painter.setBrush(self._up_brush)
|
|
else:
|
|
painter.setPen(self._down_pen)
|
|
painter.setBrush(self._down_brush)
|
|
|
|
# Draw volume body
|
|
rect: QtCore.QRectF = QtCore.QRectF(
|
|
ix - BAR_WIDTH,
|
|
0,
|
|
BAR_WIDTH * 2,
|
|
bar.volume
|
|
)
|
|
painter.drawRect(rect)
|
|
|
|
# Finish
|
|
painter.end()
|
|
return volume_picture
|
|
|
|
def boundingRect(self) -> QtCore.QRectF:
|
|
""""""
|
|
min_volume, max_volume = self._manager.get_volume_range()
|
|
rect: QtCore.QRectF = QtCore.QRectF(
|
|
0,
|
|
min_volume,
|
|
len(self._bar_picutures),
|
|
max_volume - min_volume
|
|
)
|
|
return rect
|
|
|
|
def get_y_range(self, min_ix: int | None = None, max_ix: int | None = None) -> tuple[float, float]:
|
|
"""
|
|
Get range of y-axis with given x-axis range.
|
|
|
|
If min_ix and max_ix not specified, then return range with whole data set.
|
|
"""
|
|
min_volume, max_volume = self._manager.get_volume_range(min_ix, max_ix)
|
|
return min_volume, max_volume
|
|
|
|
def get_info_text(self, ix: int) -> str:
|
|
"""
|
|
Get information text to show by cursor.
|
|
"""
|
|
bar: BarData | None = self._manager.get_bar(ix)
|
|
|
|
if bar:
|
|
text: str = f"Volume {bar.volume}"
|
|
else:
|
|
text = ""
|
|
|
|
return text
|