不同市场画图的数据单独获取

index_chart.py

@@ -0,0 +1,44 @@
+import pandas as pd
+from loguru import logger
+from chart import plot_chart, resample_data
+from db_config import DatabaseManager, DatabaseConfig
+
+
+def get_kline(code: str) -> list:
+    """
+    获取所有指数代码
+    :return:
+    """
+    env = "online"
+    env = "daily"
+    db_config = DatabaseConfig(env=env)
+    logger.info(f"数据库连接: {db_config.connection_string}")
+
+    db_manager = DatabaseManager(db_config)
+    sql = f"SELECT date as time, open, high, low, close, volume FROM public.index_kline where code='{code}' order by date;"
+    res = db_manager.execute_query(sql)
+    data_list = [dict(item) for item in res]
+    df = pd.DataFrame(data_list)
+    df["time"] = pd.to_datetime(df["time"])
+    num_cols = ["open", "high", "low", "close", "volume"]
+    for col in num_cols:
+        if col in df.columns:
+            df[col] = pd.to_numeric(df[col], errors="coerce").astype(float)
+    return df
+
+
+if __name__ == "__main__":
+    symbol = "399998"
+    timeframe = "1D"
+
+    df = pd.read_csv(
+        "/Users/aszer/Documents/vscode/etf/data/index_all_stock.csv",
+        encoding="utf-8-sig",
+    )
+    name = df.loc[df["代码"] == symbol, "名称"].values[0]
+
+    df = get_kline(code=symbol)
+    df = resample_data(df, 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)