StockAnalysisSystem/visualization.py at main · chm020924/StockAnalysisSystem · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.patches import Rectangle
from io import BytesIO
import base64
import pandas as pd

def generate_kline_chart(stock_code, stock_name, df, days=60):
    """
    生成K线图

    Returns:
        str: base64编码的图片
    """
    df = df.tail(days).copy()

    if len(df) < 10:
        return None

    try:
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(15, 10),
                                     gridspec_kw={'height_ratios': [3, 1]})

        dates = pd.to_datetime(df['日期'])
        opens = df['开盘价'] if '开盘价' in df.columns else df['收盘价']
        highs = df['最高价'] if '最高价' in df.columns else df['收盘价']
        lows = df['最低价'] if '最低价' in df.columns else df['收盘价']
        closes = df['收盘价']
        volumes = df['成交量(手)'] if '成交量(手)' in df.columns else df.get('成交量', [0]*len(df))

        for i in range(len(df)):
            date = dates.iloc[i]
            open_price = opens.iloc[i]
            high_price = highs.iloc[i]
            low_price = lows.iloc[i]
            close_price = closes.iloc[i]

            color = 'red' if close_price >= open_price else 'green'

            ax1.plot([date, date], [low_price, high_price], color='black', linewidth=0.8)
            height = abs(close_price - open_price)
            bottom = min(open_price, close_price)
            rect = Rectangle((mdates.date2num(date) - 0.3, bottom), 0.6, height,
                           facecolor=color, edgecolor='black', linewidth=0.5)
            ax1.add_patch(rect)

        if len(df) >= 5:
            ma5 = closes.rolling(window=5).mean()
            ax1.plot(dates, ma5, label='MA5', color='blue', linewidth=1)

        if len(df) >= 10:
            ma10 = closes.rolling(window=10).mean()
            ax1.plot(dates, ma10, label='MA10', color='orange', linewidth=1)

        if len(df) >= 20:
            ma20 = closes.rolling(window=20).mean()
            ax1.plot(dates, ma20, label='MA20', color='purple', linewidth=1)

        ax1.set_title(f'{stock_code} {stock_name} K线图', fontsize=16, fontweight='bold')
        ax1.set_ylabel('价格 (元)', fontsize=12)
        ax1.legend(loc='upper left')
        ax1.grid(True, alpha=0.3)

        volume_colors = ['red' if closes.iloc[i] >= opens.iloc[i] else 'green'
                       for i in range(len(df))]
        ax2.bar(dates, volumes, color=volume_colors, alpha=0.6)
        ax2.set_ylabel('成交量 (手)', fontsize=12)
        ax2.set_xlabel('日期', fontsize=12)
        ax2.grid(True, alpha=0.3)

        ax1.xaxis.set_major_formatter(mdates.DateFormatter('%m-%d'))
        ax2.xaxis.set_major_formatter(mdates.DateFormatter('%m-%d'))

        plt.setp(ax1.xaxis.get_majorticklabels(), rotation=45)
        plt.setp(ax2.xaxis.get_majorticklabels(), rotation=45)

        plt.tight_layout()

        buffer = BytesIO()
        plt.savefig(buffer, format='png', dpi=150, bbox_inches='tight')
        buffer.seek(0)
        img_base64 = base64.b64encode(buffer.read()).decode()
        plt.close()

        return img_base64

    except Exception as e:
        print(f"生成 {stock_code} K线图失败: {e}")
        plt.close()
        return None