baostock_multiprocessing 多进程取数据

#!/usr/bin/env python
import baostock as bs
import pandas as pd
import time
import os
import shutil
import multiprocessing

def download_factor(start_date, end_date, stock_df):
    rs_list = []
    result_factor = pd.DataFrame()
    for code in stock_df["code"]:
        # print("Downloading factor start:" + code,threading.current_thread().name)
        rs_factor = bs.query_adjust_factor(code=code, start_date=start_date, end_date=end_date)
        # print(rs_factor,"Downloading factor mid:" + code, threading.current_thread().name)
        while (rs_factor.error_code == '0') & rs_factor.next():
            rs_list.append(rs_factor.get_row_data())
        result_factor = pd.DataFrame(rs_list, columns=rs_factor.fields)
        # print("Downloading factor end:" + code, threading.current_thread().name)
    # print(result_factor)
    # print("Downloading factor end:" , threading.current_thread().name)
    return result_factor

def download_data(start_date,end_date,code):
    # 获取指定日期的指数、股票数据
    data_df = pd.DataFrame()
    #print("Downloading :" + code)
    k_rs = bs.query_history_k_data_plus(code, "date,code,open,high,low,close,volume,amount,turn,pctChg,peTTM,pbMRQ,psTTM,pcfNcfTTM",
                                        start_date=start_date, end_date=end_date,adjustflag= "2",frequency="d")
    data_df = data_df.append(k_rs.get_data())
    return data_df

def conpare_list():
    stock_rs = bs.query_all_stock(end_date)
    stock_df = stock_rs.get_data()
    file_name = pathsave + "\\" + "all.csv"
    print(file_name)
    stock_read = pd.read_csv(file_name)
    print(stock_read)
    for code in stock_df["code"]:
        #print(code)
        flag_t = stock_read.loc[stock_read["code"] == code,"flag"]
        flag_t = flag_t.reset_index(drop=True)
        flag_t = pd.DataFrame(flag_t)
        t = ''
        if flag_t.empty:
            t = "new"
        else:
            t = flag_t.loc[0,"flag"]
        stock_df.loc[stock_df["code"] == code,"flag"] = t
    return stock_df

def add_data(end_date,stock_df,pathsave):
    stock_df = stock_df.drop_duplicates(subset=["code"], keep="last", inplace=False)
    stock_df["code2"] = stock_df["code"].str.replace("sh.", "SH")
    stock_df["code2"] = stock_df["code2"].str.replace("sz.", "SZ")
    stock_df = stock_df.set_index("code")
    #print(stock_df)
    for code in stock_df.index:
        file = pathsave + "\\"  + stock_df.loc[code,"flag"]  +"\\"+ stock_df.loc[code,"code2"]+".csv"
        #print(file)
        df_old = pd.DataFrame()
        if  os.path.isfile(file):
            df_old = pd.read_csv(file)
        df_all = download_data(stock_df.loc[code,"start_date"],end_date,code)
        df_all["code"] = df_all["code"].str.replace("sh.", "SH")
        df_all["code"] = df_all["code"].str.replace("sz.", "SZ")
        df_all["date"] = df_all["date"].str.replace("-", "")
        df_old = df_old.append(df_all)
        #df_new = df_old.reset_index(drop=True)
        df_old["date"] = df_old["date"].astype(str)
        df_old = df_old.drop_duplicates(subset=["date"], keep="last", inplace=False)
        df_old.to_csv(file,sep=",",encoding="gbk", index=False)

def rewrite_new_file(pathsave):#对新增加的股票进行移动，更新到all.csv文件
    file_name_w = pathsave + "\\" + "all.csv"
    file_name_r = pathsave + "\\" + "list.csv"
    pathdir = pathsave + "\\" + "new"
    stock_read = pd.read_csv(file_name_r)
    pd_new = stock_read.loc[stock_read["flag"] == "new"]
    #newfiles = os.listdir(pathdir)
    #print(stock_read)
    if len(pd_new)>0:
        for file1 in pd_new["code"]:
            file =file1
            #print(file)
            file = file.replace("sz.", "SZ")
            file = file.replace("sh.", "SH")
            file = file + ".csv"
            file2 = file
            file = pathdir + "\\" + file
            if os.path.isfile(file):
                df_new = pd.read_csv(file)
                if pd.isna(df_new.loc[0,"peTTM"]):
                    print(file,"可能是指数文件")
                else:
                    if file.find("SZ")>=0:
                        #print(file.find("SZ"))
                        stock_read.loc[stock_read["code"]==file1, "flag"] = "sz"
                        pathdir_sz = pathsave + "\\" + "sz"
                        dstfile = pathdir_sz +"\\"+file2
                        shutil.move(file, dstfile)
                    else:
                        stock_read.loc[stock_read["code"]==file1, "flag"] = "sh"
                        pathdir_sz = pathsave + "\\" + "sh"
                        dstfile = pathdir_sz + "\\" + file2
                        shutil.move(file, dstfile)

    stock_read.to_csv(file_name_w,sep=",",encoding="utf-8", index=False)

def sub_process(start_date,end_date,df_only_name1,q):
    lg = bs.login()
    print('login respond error_code:' + lg.error_code)
    print('login respond  error_msg:' + lg.error_msg)
    print('-----process begin-----')
    print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), multiprocessing.current_process().name)
    df_factor1 = download_factor(start_date, end_date, df_only_name1)
    q.put(df_factor1,block = False)
    print('-----process done-----')
    print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),multiprocessing.current_process().name)
    exit(0)

def sub_process2(end_date,df_only_name1,pathsave,q):
    lg = bs.login()
    print('login respond error_code:' + lg.error_code)
    print('login respond  error_msg:' + lg.error_msg)
    print('-----process 下载数据 begin-----')
    print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), multiprocessing.current_process().name)
    add_data(end_date, df_only_name1,pathsave)
    q.put(multiprocessing.current_process().name,block = False)
    print('-----process 数据下载写入结束 done-----')
    print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),multiprocessing.current_process().name)
    exit(0)

if __name__ == '__main__':
    # 获取指定日期全部股票的日K线数据
    print("hello")
    lg = bs.login()
    print('login respond error_code:' + lg.error_code)
    print('login respond  error_msg:' + lg.error_msg)
    pathsave = 'G:\\datas of status\\python codes\\baostock\\lx'  # 设定临时文件存放位置

    ori_date = "2018-01-01"#设定最初日期数据
    start_date = "2020-08-18"     #常设，设定这次要下载的数据开始日期
    end_date = "2020-08-20"       #常设，设定这次要下载的数据结束日期,结束日期必须是交易日，否则会出错
    print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
    print("开始比较")
    stock_df = conpare_list()    #分清指数，上证，深证
    print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
    print("开始下载factor")
    file_w = pathsave + "\\" + "list.csv"
    stock_df.to_csv(file_w, sep=",", index=False, header=True)

    #=====================下载factor
    all_nums = len(stock_df)
    epochs = 5
    step = int(all_nums / epochs)
    process_list = []
    q = multiprocessing.Queue(maxsize=epochs)
    for i in range(epochs):
        begin = i * step
        end = begin + step
        if i == epochs - 1:
            end = all_nums
        df_only_name1 = stock_df[begin:end]
        print("no.",i,begin,end)
        tmp_process = multiprocessing.Process(target=sub_process, args=(start_date,end_date,df_only_name1, q))
        process_list.append(tmp_process)
    for process in process_list:
        process.start()
        # print("start",process)
    while (q.qsize() != epochs):
        # print(q.qsize(),"begin")
        if (q.qsize() >= 1):
            print(q.qsize())
            time.sleep(5)
        else:
            time.sleep(20)

    time.sleep(1)
    df_factor = pd.DataFrame()
    while not q.empty():
        list_g = q.get()
        df_factor = df_factor.append(list_g)
    #=========
    #df_factor = download_factor(start_date,end_date,stock_df)  #分清有无复权，若有则设定开初下载数据时间有最初日期，然后再重新下载数据
    df_factor = df_factor.drop_duplicates(subset=["code"], keep="last", inplace=False)
    print(df_factor)
    print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
    #exit(0)

    print("下载factor结束，开始下载数据")
    stock_df["start_date"] = start_date
    for code in df_factor["code"]:
        stock_df.loc[stock_df["code"] == code,"start_date"] = ori_date
    #print(stock_df[220:240])
    print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),"下边开始下载数据")
    #==============================#下载数据
    all_nums = len(stock_df)
    epochs = 5
    step = int(all_nums / epochs)
    process_list = []
    q = multiprocessing.Queue(maxsize=epochs)
    for i in range(epochs):
        begin = i * step
        end = begin + step
        if i == epochs - 1:
            end = all_nums
        df_only_name1 = stock_df[begin:end]
        print("no.", i, begin, end)
        tmp_process = multiprocessing.Process(target=sub_process2, args=(end_date, df_only_name1,pathsave, q))
        process_list.append(tmp_process)
    for process in process_list:
        process.start()
        # print("start",process)
    while (q.qsize() != epochs):
        # print(q.qsize(),"begin")
        if (q.qsize() >= 1):
            print(q.qsize())
            time.sleep(5)
        else:
            time.sleep(20)

    time.sleep(1)
    #df_process = pd.DataFrame()
    while not q.empty():
        list_g = q.get()
        print(list_g,"done")
        #df_process = df_process.append(list_g)
    #=============================
    print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
    print("下载数据结束")
    rewrite_new_file(pathsave)
    #print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
    bs.logout()