import csv
import time
from selenium import webdriver

MAX_RETRY_TIME = 150
URL = "https://emweb.securities.eastmoney.com/pc_hsf10/pages/index.html?type=web&code=%s&color=b#/gsgk"

def get_company(stock_code_list):
    try:
        browser = webdriver.Firefox()
        for stock_code in stock_code_list:
            code = stock_code[0]
            name = stock_code[1]
            stock_code = stock_code[0]
            if stock_code[0] == '6':
                stock_code = 'SH' + stock_code
            elif stock_code[0] == '8' or stock_code[0] == '4':
                stock_code = 'BJ' + stock_code
            else:
                stock_code = 'SZ' + stock_code
            url = URL % stock_code
            browser.execute_script('window.open("%s")' % url)

            retry_time = 0
            while True:
                try:
                    if retry_time < MAX_RETRY_TIME:
                        browser.switch_to.window(browser.window_handles[1])
                        company_name = browser.find_elements_by_css_selector('div.jbzl_table tr:nth-child(1) td')[0].text
                        break
                    else:
                        return "获取失败"
                except:
                    retry_time += 1
                    time.sleep(0.2)

            company_details = {"股票代码":code,"股票名称":name,"公司名称":company_name}
            print(company_details)
            company_details_csv = [code,name,company_name]
            with open('code_2.csv', 'a', newline='') as newfile:
                writer = csv.writer(newfile)
                writer.writerow(company_details_csv)
                newfile.close()
            browser.close()
            browser.switch_to.window(browser.window_handles[0])

        browser.close()

    except Exception as e:
        print(e)
        return "出错啦"

file = 'stock_code.csv'
with open(file,'r') as file:
    code_list = csv.reader(file,delimiter=',')
    stock_code_list = []
    for code in code_list:
        stock_code_list.append(code)
    del stock_code_list[0]
    get_company(stock_code_list)