VeighNa量化社区
你的开源社区量化交易平台
Member
avatar
加入于:
帖子: 3
声望: 1

哈,感谢楼主,顺便做了点改进,把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"
Member
avatar
加入于:
帖子: 33
声望: 0

star990 wrote:

哈,感谢楼主,顺便做了点改进,把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"

大佬,能贴给清洁版吗?

Member
avatar
加入于:
帖子: 3
声望: 0

在vnpy 3.2.0版本下要如何修改代码才能跑起来?

Member
avatar
加入于:
帖子: 8
声望: 0

不错,数据可视化很重要

© 2015-2022 上海韦纳软件科技有限公司
备案服务号:沪ICP备18006526号

沪公网安备 31011502017034号

【用户协议】
【隐私政策】
【免责条款】