根据period随机生成颜色；添加买卖信号markers

chart.py

@@ -5,18 +5,32 @@ import talib as ta
 from lightweight_charts import Chart
 
 
+import random
+
+
+def get_fixed_color_based_on_period(period):
+    # 使用周期值作为随机种子，确保相同周期生成相同颜色
+    random.seed(period)
+
+    # 生成随机的RGB值
+    r = random.randint(0, 255)
+    g = random.randint(0, 255)
+    b = random.randint(0, 255)
+
+    # 将RGB值转换为十六进制颜色代码
+    color_code = "#{:02x}{:02x}{:02x}".format(r, g, b)
+
+    # 重置随机种子，避免影响其他随机操作
+    random.seed(None)
+
+    return color_code
+
+
 def add_ema(df, chart, period: int = 50):
-    ema = ta.EMA(df["close"], timeperiod=period)
     name = f"EMA {period}"
-    df[name] = ema
-    PERIOD_COLORS = {
-        10: "#E42C2A",
-        20: "#E49D30",
-        30: "#3CFF0A",
-        60: "#3275E4",
-        200: "#B06CE3",
-    }
-    line = chart.create_line(name, color=PERIOD_COLORS[period], width=2)
+    df[name] = ta.EMA(df["close"], timeperiod=period)
+    color = get_fixed_color_based_on_period(period)
+    line = chart.create_line(name, color=color, width=2)
     line.set(df[["time", name]])
 
 
@@ -114,6 +128,38 @@ def add_TD(df, chart):
     chart.marker_list(markers)
 
 
+def add_buy_sell_signal_markers(df, chart):
+
+    if "buy" not in df.columns or "sell" not in df.columns:
+        return
+
+    markers = []
+    signals = df[["time", "buy", "sell"]].to_dict(orient="records")
+    for item in signals:
+        if item["buy"] == 1:
+            markers.append(
+                {
+                    "time": item["time"].strftime("%Y-%m-%d"),
+                    "position": "below",
+                    "shape": "arrow_up",
+                    "color": "#4CAF50",
+                    "text": f"B",
+                }
+            )
+        elif item["sell"] == 1:
+            markers.append(
+                {
+                    "time": item["time"].strftime("%Y-%m-%d"),
+                    "position": "above",
+                    "shape": "arrow_down",
+                    "color": "#F44336",
+                    "text": f"S",
+                }
+            )
+
+    chart.marker_list(markers)
+
+
 def get_kline(code: str) -> list:
     """
     获取所有指数代码
@@ -179,7 +225,12 @@ def resample_data(df: pd.DataFrame, timeframe: str) -> pd.DataFrame:
     return resampled
 
 
-def plot_chart(symbol: str, name: str, timeframe: str):
+def plot_chart(df, symbol: str, name: str, timeframe: str):
+    """
+    df: DataFrame
+    time, open, high, low, close, volume, [buy, sell]可选 buy=1 买入信号, sell=1 卖出信号
+
+    """
     chart = Chart(toolbox=True, inner_height=0.7, maximize=True)
 
     chart.topbar.textbox("symbol", symbol)
@@ -196,6 +247,7 @@ def plot_chart(symbol: str, name: str, timeframe: str):
     add_macd(df, chart)
     add_cci(df, chart, period=14)
     add_TD(df, chart)
+    add_buy_sell_signal_markers(df, chart)
 
     chart.show(block=True)
 
@@ -212,4 +264,6 @@ if __name__ == "__main__":
 
     df = get_kline(code=symbol)
     df = resample_data(df, timeframe)
-    plot_chart(symbol=symbol, name=name, timeframe=timeframe)
+    # df['buy'] = df['time'].apply(lambda x: 1 if pd.to_datetime(x).day % 2 == 1 else 0)
+    # df['sell'] = df['time'].apply(lambda x: 1 if pd.to_datetime(x).day % 2 == 0 else 0)
+    plot_chart(df=df, symbol=symbol, name=name, timeframe=timeframe)