# 加載取數(shù)與繪圖所需的函數(shù)包
import pandas as pd
import datetime
from hs_udata import set_token,stock_quote_daily
from mpl_finance import candlestick_ohlc
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
mpl.rcParams['font.sans-serif'] = ['SimHei'] # 指定默認字體
mpl.rcParams['axes.unicode_minus'] = False  # 解決保存圖像是負號'-'顯示為方塊的問題

def GetData(stock_code,start,end):
    #stock_code：獲取股票數(shù)據(jù)的股票代碼
    #     start：開始日期
    #       end：結(jié)束日期
    date_start=datetime.datetime.strptime(start,'%Y-%m-%d')
    date_end  =datetime.datetime.strptime(end,'%Y-%m-%d')
    data = pd.DataFrame([])
    while date_start<date_end:
        # 獲取日行情數(shù)據(jù)，接口說明見 https://udata.hs.net/datas/332/
        # adjust_way枚舉值為：0-不復權(quán)，1-前復權(quán)，2-后復權(quán)，此處取前復權(quán)
        data_i = stock_quote_daily(en_prod_code=stock_code
                                   ,trading_date=date_start.strftime('%Y%m%d')
                                   ,adjust_way = 1)
        data=pd.concat([data,data_i],axis=0)      # 將行情數(shù)據(jù)按行拼接
        date_start+=datetime.timedelta(days=1)    # 日期變量自增
    # 返回行情數(shù)據(jù)
    return data

#1、獲取行情數(shù)據(jù)
stock_code = "600570.SH"     # 恒生電子 股票代碼是600570.SH
start='2021-06-01'
end  ='2021-08-01'
set_token(token = 'xxxxxxxxxxxxxxxxxxxxxxxx')   # 注冊后，獲取并替換token
data = GetData(stock_code,start,end)
data

#2、數(shù)據(jù)處理
data = data.loc[data.turnover_status=='交易']             # 剔除非交易日
data_price = data[['trading_date','open_price','high_price','low_price','close_price'
                   ,'business_amount']]                  # 選取日期、高開低收價格、成交量數(shù)據(jù)
data_price.set_index('trading_date', inplace=True)       # 將日期作為索引
data_price = data_price.astype(float)                    # 將價格數(shù)據(jù)類型轉(zhuǎn)為浮點數(shù)
# 將日期格式轉(zhuǎn)為 candlestick_ohlc 可識別的數(shù)值
data_price['Date'] = list(map(lambda x:mdates.date2num(datetime.datetime.strptime(x,'%Y-%m-%d')),data_price.index.tolist()))
data_price

#3、繪制K線圖
# 提取繪圖數(shù)據(jù)
ohlc = data_price[['Date','open_price','high_price','low_price','close_price']]
f1, ax = plt.subplots(figsize = (12,6))                        # 創(chuàng)建圖片
candlestick_ohlc(ax, ohlc.values.tolist(), width=.7
                 , colorup='red', colordown='green')           # 使用candlestick_ohlc繪圖
ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d')) # 設置橫軸日期格式
plt.xticks(rotation=30)                                        # 日期顯示的旋轉(zhuǎn)角度
plt.title(stock_code,fontsize = 14)                            # 設置圖片標題
plt.xlabel('日 期',fontsize = 14)                              # 設置橫軸標題
plt.ylabel('價 格（元）',fontsize = 14)                        # 設置縱軸標題
plt.show()

# 4、去除非交易日的間隔
ohlc = data_price[['Date','open_price','high_price','low_price','close_price']]
ohlc.loc[:,'Date'] = range(len(ohlc))     # 重新賦值橫軸數(shù)據(jù)，使橫軸數(shù)據(jù)為連續(xù)數(shù)值
# 繪圖
f1, ax = plt.subplots(figsize = (12,6))
candlestick_ohlc(ax, ohlc.values.tolist(), width=.7, colorup='red', colordown='green')
plt.xticks(rotation=30)                                        # 日期顯示的旋轉(zhuǎn)角度
plt.title(stock_code,fontsize = 14)                            # 設置圖片標題
plt.xlabel('日 期',fontsize = 14)                               # 設置橫軸標題
plt.ylabel('價 格（元）',fontsize = 14)                          # 設置縱軸標題
# 修改橫軸標注日期
date_list = ohlc.index.tolist()           # 獲取日期列表
xticks_len = round(len(date_list)/(len(ax.get_xticks())-1))    # 獲取默認橫軸標注的間隔
xticks_num = range(0,len(date_list),xticks_len)                # 生成橫軸標注位置列表
xticks_str = list(map(lambda x:date_list[int(x)],xticks_num))  # 生成正在標注日期列表
ax.set_xticks(xticks_num)                                      # 設置橫軸標注位置
ax.set_xticklabels(xticks_str)                                 # 設置橫軸標注日期
plt.show()

#5、繪制成交量
fig = plt.figure(figsize=(12,10))
grid = plt.GridSpec(12, 10, wspace=0.5, hspace=0.5)
#（1）繪制K線圖
# K線數(shù)據(jù)
ohlc = data_price[['Date','open_price','high_price','low_price','close_price']]
ohlc.loc[:,'Date'] = range(len(ohlc))     # 重新賦值橫軸數(shù)據(jù)，繪制K線圖無間隔
# 繪制K線
ax1 = fig.add_subplot(grid[0:8,0:12])   # 設置K線圖的尺寸
candlestick_ohlc(ax1, ohlc.values.tolist(), width=.7
                 , colorup='red', colordown='green')
plt.title(stock_code,fontsize = 14)     # 設置圖片標題
plt.ylabel('價 格（元）',fontsize = 14)   # 設置縱軸標題
ax1.set_xticks([])                      # 日期標注在成交量中，故清空此處x軸刻度
ax1.set_xticklabels([])                 # 日期標注在成交量中，故清空此處x軸 
#（2）繪制成交量
# 成交量數(shù)據(jù)
data_volume = data_price[['Date','close_price','open_price','business_amount']]
data_volume['color'] = data_volume.apply(lambda row: 1 if row['close_price'] >= row['open_price'] else 0, axis=1)        # 計算成交量柱狀圖對應的顏色，使之與K線顏色一致
data_volume.Date = ohlc.Date
# 繪制成交量
ax2 = fig.add_subplot(grid[8:10,0:12])  # 設置成交量圖形尺寸
ax2.bar(data_volume.query('color==1')['Date']
        , data_volume.query('color==1')['business_amount']
        , color='r')                    # 繪制紅色柱狀圖
ax2.bar(data_volume.query('color==0')['Date']
        , data_volume.query('color==0')['business_amount']
        , color='g')                    # 繪制綠色柱狀圖
plt.xticks(rotation=30) 
plt.xlabel('日 期',fontsize = 14)                               # 設置橫軸標題
# 修改橫軸日期標注
date_list = ohlc.index.tolist()           # 獲取日期列表
xticks_len = round(len(date_list)/(len(ax2.get_xticks())-1))      # 獲取默認橫軸標注的間隔
xticks_num = range(0,len(date_list),xticks_len)                   # 生成橫軸標注位置列表
xticks_str = list(map(lambda x:date_list[int(x)],xticks_num))     # 生成正在標注日期列表
ax2.set_xticks(xticks_num)                                        # 設置橫軸標注位置
ax2.set_xticklabels(xticks_str)                                   # 設置橫軸標注日期
plt.show()