import baostock as bs
import pandas as pd
import talib


def code_to_symbol(code):
    return 'sh.%s' % code if code[:1] in ['5', '6', '9'] or code[:2] in ['11', '13'] else 'sz.%s' % code


def query_k_data(code, start, end, frequency='d'):
    bs.login()
    codeSymbol = code_to_symbol(code)
    rs = bs.query_history_k_data_plus(codeSymbol,
                                      "date,open,high,low,close,preclose,volume,amount,turn,tradestatus,pctChg",
                                      start_date=start, end_date=end,
                                      frequency=frequency, adjustflag="2")
    print(rs)
    data_list = []
    while (rs.error_code == '0') & rs.next():
        data_list.append(rs.get_row_data())
    df = pd.DataFrame(data_list, columns=rs.fields)
    df = df.loc[df.tradestatus == '1']
    df.open = df.open.astype(float)
    df.high = df.high.astype(float)
    df.low = df.low.astype(float)
    df.close = df.close.astype(float)
    bs.logout()
    return df


def fill_macd(df):
    macd, macdsignal, macdhist = talib.MACD(df['close'].values)
    df['macd'] = macd
    return df


if __name__ == '__main__':
    df = query_k_data('000970', '2018-01-01', '2018-05-16')
    print(fill_macd(df))