python计算指定股票分红及回报率

import baostock as bs
import pandas as pd
from datetime import datetime

class Stock:
    def __init__(self):
        self.bs=bs
        self.lg=bs.login()
        #定义股票代码
        self.stockcode='sh.601288'
        #定义定投金额
        self.by_amount=3000
        #定义持仓股数
        self.stocknum=0
        #定义持仓金额
        self.have_amount=0
        #定义总投入金额
        self.total_amount=0
        #定义购买时间
        self.by_date='07-01'
        #定义开始时间
        self.start_date='2010-01-01'
        #定义截止时间
        self.end_date='2021-12-31'
        #定义余额
        self.amount=0
        #定义收益率
        self.get_rate=0
        #定义分红总金额
        self.divid_amount=0
        #定义是否红利再投
        self.all_by=True
        #定义投资表
        self.result=[]

    #获取指定日期区间的年份数据
    def get_years(self,start_date_str, end_date_str):
        start_date = datetime.strptime(start_date_str, "%Y-%m-%d")
        end_date = datetime.strptime(end_date_str, "%Y-%m-%d")
        start_year = start_date.year
        end_year = end_date.year
        years = list(range(start_year, end_year + 1))
        return years

    #获取指定股票的分红除权信息
    def get_divid_info(self,stockcode,start_date,end_date):
        years=self.get_years(start_date, end_date)
        rs_list = []
        for i in range(0,len(years)):
            rs_dividend = bs.query_dividend_data(code=stockcode, year=years[i], yearType="operate")
            while (rs_dividend.error_code == '0') & rs_dividend.next():
                rs_list.append(rs_dividend.get_row_data())
        return rs_list
    #获取指定股票的历史日K数据
    def get_stock_data(self,stockcode,start_date,end_date):
        rs = bs.query_history_k_data_plus(stockcode,
            "date,code,open,high,low,close",
        start_date=start_date, end_date=end_date,
        frequency="d", adjustflag="3")
        data_list = []
        while (rs.error_code == '0') & rs.next():
            data_list.append(rs.get_row_data())
        return data_list
    #计算定投购买数据
    def by_stock(self,stock_data,divid_data):
        needby=False
        fh=0
        spj=0
        for item in stock_data:
            if needby==True:
                one_row=[]
                #红利不投
                if self.all_by==False:
                    self.divid_amount=self.divid_amount+fh*self.stocknum #分红总金额
                    gj=float(item[5])#当前股价
                    grnum=int((self.by_amount+self.amount)/gj) #购入股数
                    self.stocknum=self.stocknum+grnum#购入整数股票
                    self.total_amount=self.total_amount+self.by_amount #总投入
                    self.amount=self.amount+self.by_amount-(gj*grnum) #余额
                    needby=False
                    one_row.append(item[0])#加入日期
                    one_row.append(gj)#加入股价
                    one_row.append(self.total_amount)#加入投资总额
                    one_row.append(self.stocknum)#加入持有股票数量
                    one_row.append(round(self.stocknum*gj,2))#加入持仓总额
                    one_row.append(round(self.amount,2))#加入余额
                    one_row.append(fh)#加入分红股利
                    one_row.append(round(fh*self.stocknum,2))#加入分红
                    self.get_rate=(((self.divid_amount+(self.stocknum*gj)+self.amount)-self.total_amount)/self.total_amount)
                    one_row.append(round(self.get_rate*100,2))#加入收益率
                    self.result.append(one_row)

                else:
                    #红利再投
                    self.divid_amount=self.divid_amount+fh*self.stocknum #分红总金额
                    divid_amount=fh*self.stocknum #分红金额
                    gj=float(item[5])#当前股价
                    grnum=int((self.by_amount+self.amount+divid_amount)/gj) #购入股数
                    self.stocknum=self.stocknum+grnum#购入整数股票
                    self.total_amount=self.total_amount+self.by_amount #总投入
                    self.amount=self.amount+divid_amount+self.by_amount-(gj*grnum) #余额
                    needby=False
                    one_row.append(item[0])#加入日期
                    one_row.append(gj)#加入股价
                    one_row.append(self.total_amount)#加入投资总额
                    one_row.append(self.stocknum)#加入持有股票数量
                    one_row.append(round(self.stocknum*gj,2))#加入持仓总额
                    one_row.append(round(self.amount,2))#加入余额
                    one_row.append(fh)#加入分红股利
                    one_row.append(round(fh*self.stocknum,2))#加入分红
                    self.get_rate=((((self.stocknum*gj)+self.amount)-self.total_amount)/self.total_amount)
                    one_row.append(round(self.get_rate*100,2))#加入收益率
                    self.result.append(one_row)
            for oned in divid_data:
                if oned[7]== item[0]:
                    fh=float(oned[9])#每股分红
                    spj=float(item[5])#收盘价
                    needby=True
            
        return self.result

#测试
if __name__ == '__main__':
    stk=Stock()
    #stk.all_by=False
    stk.by_amount=30000
    stk.stockcode='sh.601939'
    divid_data=stk.get_divid_info(stk.stockcode,"2011-01-01","2021-12-31")
    stock_data=stk.get_stock_data(stk.stockcode,"2010-01-01","2021-12-31")
    rs=stk.by_stock(stock_data,divid_data)
    print(rs)
    result = pd.DataFrame(rs, columns=['日期','股价','投入','股数','持仓','余额','股利','分红','收益率'])

    result.to_csv("C:\\Users\daimadog\Desktop\建设银行.csv", index=False)