自己试着改了下,用了一个二维字典存储各个周期的均线图,目前运行还可以(有点小瑕疵),把代码po出来,就当抛砖引玉了。
class SmaItem(CandleItem):
""""""
def __init__(self, manager: BarManager):
""""""
super().__init__(manager)
self.blue_pen: QtGui.QPen
#self.sma_window = 10
#self.sma_data: Dict[int, float] = {} #modified by lnt
self.sma_data : Dict = {}
#二维数据结构
def addtwodimdict(self,thedict, key_a, key_b, val) -> dict:
if key_a in thedict:
thedict[key_a].update({key_b: val})
else:
thedict.update({key_a:{key_b: val}})
return thedict
#def get_sma_value(self, ix: int) -> float:
def get_sma_value(self, ix: int,ma_sd:list) -> list:#modified by lnt
""""""
ma_rt = []
if ix < 0:
for i in range(len(ma_sd)):#全部赋值nan
ma_rt.append(math.nan)
return ma_rt
# When initialize, calculate all rsi value
if not self.sma_data:
bars = self._manager.get_all_bars()
close_data = [bar.close_price for bar in bars]
for m in range(len(ma_sd)):
sma_array = talib.SMA(np.array(close_data), ma_sd[m])
for n, value in enumerate(sma_array):
self.sma_data = self.addtwodimdict(self.sma_data,ma_sd[m],n,value)
# Return if already calcualted
if ix in self.sma_data[ma_sd[0]]:
for i in range(len(ma_sd)):
ma_rt.append(self.sma_data[ma_sd[i]][ix])
return ma_rt
# Else calculate new value
close_data = []
for n in range(ix - self.sma_window, ix + 1):
bar = self._manager.get_bar(n)
close_data.append(bar.close_price)
for n in range(len(ma_sd)):
sma_array = talib.SMA(np.array(close_data), ma_sd[n])
sma_value = sma_array[-1]
self.sma_data = self.addtwodimdict(self.sma_data,ma_sd[n],ix,sma_value)
ma_rt.append(sma_value)
return ma_rt
def _draw_bar_picture(self, ix: int, bar: BarData) -> QtGui.QPicture:
""""""
ma = [5,10,20,30,60,120]
sma_value = self.get_sma_value(ix,ma)
last_sma_value = self.get_sma_value(ix - 1,ma)
# Create objects
picture = QtGui.QPicture()
painter = QtGui.QPainter(picture)
# Set painter color
pencolor =[(255, 0, 0),(255, 255, 255),(255, 100, 255),(255, 255, 0),(0,255,0),(0,0,255)]
penwidth = [1,1,1,2,1,1,1,1]
# Draw Line
for k in range(len(ma)):
start_point = QtCore.QPointF(ix-1, last_sma_value[k])
end_point = QtCore.QPointF(ix, sma_value[k])
self.blue_pen = pg.mkPen(color=pencolor[k], width=penwidth[k])
painter.setPen(self.blue_pen)
painter.drawLine(start_point, end_point)
# Finish
painter.end()
return picture