哈,感谢楼主,顺便做了点改进,把ZB/ZB2归为一个class,diff如下:
diff --git a/chart/widget.py b/chart/widget.py
index 8d6ff34..c852d32 100755
--- a/chart/widget.py
+++ b/chart/widget.py
@@ -4,6 +4,7 @@ import pyqtgraph as pg
from vnpy.trader.ui import QtGui, QtWidgets, QtCore
from vnpy.trader.object import BarData
+from vnpy.trader.constant import Color
from .manager import BarManager
from .base import (
@@ -114,12 +115,17 @@ class ChartWidget(pg.PlotWidget):
self,
item_class: Type[ChartItem],
item_name: str,
- plot_name: str
+ plot_name: str,
+ color: "Color" = None,
+ line: str = None
):
"""
Add chart item.
"""
- item = item_class(self._manager)
+ if color is not None and line is not None:
+ item = item_class(self._manager, color.value, line)
+ else:
+ item = item_class(self._manager)
self._items[item_name] = item
plot = self._plots.get(plot_name)
diff --git a/examples/chart_draw/demo.py b/examples/chart_draw/demo.py
index 5c0a4d3..b2adc51 100644
--- a/examples/chart_draw/demo.py
+++ b/examples/chart_draw/demo.py
@@ -1,6 +1,6 @@
from datetime import datetime
from typing import Dict
-from vnpy.trader.constant import Exchange, Interval
+from vnpy.trader.constant import Exchange, Interval, Color
from vnpy.trader.database import database_manager
from vnpy.app.cta_strategy import ArrayManager
from vnpy.chart import ChartWidget, VolumeItem, CandleItem
@@ -9,16 +9,16 @@ from vnpy.trader.ui import create_qapp, QtCore, QtGui
from vnpy.trader.object import BarData
from vnpy.chart.manager import BarManager
-
class ZB(CandleItem):
"""自定义指标显示"""
- def __init__(self, manager: BarManager):
+ def __init__(self, manager: BarManager, color: "Color", line: str):
""""""
super().__init__(manager)
- self.blue_pen: QtGui.QPen = pg.mkPen(color=(100, 100, 255), width=2)
+ self.pen: QtGui.QPen = pg.mkPen(color, width=2)
self.sma_data: Dict[int, float] = {}
+ self.line = line
def get_sma_value(self, ix: int) -> float:
""""""
@@ -27,7 +27,7 @@ class ZB(CandleItem):
if not self.sma_data:
bars = self._manager.get_all_bars()
- sma_array = [bar.down_line for bar in bars]
+ sma_array = [bar.__dict__[self.line] for bar in bars]
for n, value in enumerate(sma_array):
self.sma_data[n] = value
@@ -49,7 +49,7 @@ class ZB(CandleItem):
painter = QtGui.QPainter(picture)
# Set painter color
- painter.setPen(self.blue_pen)
+ painter.setPen(self.pen)
# Draw Line
start_point = QtCore.QPointF(ix-1, last_sma_value)
@@ -70,67 +70,6 @@ class ZB(CandleItem):
return text
-class ZB2(CandleItem):
- """自定义指标显示"""
-
- def __init__(self, manager: BarManager):
- """"""
- super().__init__(manager)
-
- self.blue_pen: QtGui.QPen = pg.mkPen(color=(100, 100, 255), width=2)
- self.sma_data: Dict[int, float] = {}
-
- def get_sma_value(self, ix: int) -> float:
- """"""
- if ix < 0:
- return 0
-
- if not self.sma_data:
- bars = self._manager.get_all_bars()
- sma_array = [bar.up_line for bar in bars]
-
- for n, value in enumerate(sma_array):
- self.sma_data[n] = value
-
- if ix in self.sma_data:
- return self.sma_data[ix]
-
- sma_value = sma_array[-1]
-
- return sma_value
-
- def _draw_bar_picture(self, ix: int, bar: BarData) -> QtGui.QPicture:
- """"""
- sma_value = self.get_sma_value(ix)
- last_sma_value = self.get_sma_value(ix - 1)
-
- # Create objects
- picture = QtGui.QPicture()
- painter = QtGui.QPainter(picture)
-
- # Set painter color
- painter.setPen(self.blue_pen)
-
- # Draw Line
- start_point = QtCore.QPointF(ix-1, last_sma_value)
- end_point = QtCore.QPointF(ix, sma_value)
- painter.drawLine(start_point, end_point)
-
- # Finish
- painter.end()
- return picture
-
- def get_info_text(self, ix: int) -> str:
- """"""
- if ix in self.sma_data:
- sma_value = self.sma_data[ix]
- text = f"ZB {sma_value:.2f}"
- else:
- text = "ZB -"
-
- return text
-
-
if __name__ == "__main__":
app = create_qapp()
@@ -175,8 +114,8 @@ if __name__ == "__main__":
widget.add_item(CandleItem, "candle", "candle")
widget.add_item(VolumeItem, "volume", "volume")
- widget.add_item(ZB, "ZB", "candle")
- widget.add_item(ZB2, "ZB2", "candle")
+ widget.add_item(ZB, "ZB", "candle", Color.Blue, "down_line")
+ widget.add_item(ZB, "ZB2", "candle", Color.Green, "up_line")
widget.add_cursor()
history = bars
diff --git a/trader/constant.py b/trader/constant.py
index 2612cb5..995e6d4 100755
--- a/trader/constant.py
+++ b/trader/constant.py
@@ -157,3 +157,12 @@ class Interval(Enum):
DAILY = "d"
WEEKLY = "w"
TICK = "tick"
+
+class Color(Enum):
+ """
+ Color.
+ """
+ Red = "#FF0000"
+ Green = "#00FF00"
+ Blue = "#0000FF"
+ White = "#FFFFFF"
支持海韦纳软件正版vn.py, 支持陈晓优,强列谴责上海量贝+量化林+王登高的下作行为