diff --git a/chart.py b/chart.py index 6314f02..cffd920 100644 --- a/chart.py +++ b/chart.py @@ -1,5 +1,5 @@ import pandas as pd -from db_config import DatabaseManager, DatabaseConfig + from loguru import logger import talib as ta import numpy as np @@ -30,7 +30,7 @@ def get_fixed_color_based_on_period(period): def add_ema(df, chart, period: int = 50): - name = f"EMA {period}" + name = f"EMA_{period}" df[name] = ta.EMA(df["close"], timeperiod=period) color = get_fixed_color_based_on_period(period) line = chart.create_line( @@ -206,29 +206,6 @@ def add_buy_sell_signal_markers(df, chart): chart.marker_list(markers) -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 - - def resample_data(df: pd.DataFrame, timeframe: str) -> pd.DataFrame: """ 对日线数据进行重采样 @@ -372,10 +349,10 @@ def plot_chart(df, symbol: str, name: str, timeframe: str): chart.set(df) chart.events.range_change += on_range_change_poc - add_ema(df, chart, period=10) - add_ema(df, chart, period=20) + # add_ema(df, chart, period=10) + # add_ema(df, chart, period=20) add_ema(df, chart, period=30) - add_ema(df, chart, period=60) + # add_ema(df, chart, period=60) add_macd(df, chart) add_cci(df, chart, period=14) add_TD(df, chart) @@ -385,17 +362,4 @@ def plot_chart(df, symbol: str, name: str, timeframe: str): if __name__ == "__main__": - symbol = "399998" - timeframe = "1W" - - 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) + ...