import os

def fetch(data):  # 查询功能函数

    ret = ''  # 查询结果的返回列表

    tag = False

    with open('沪市股票', 'r') as hu_stock:

        for read_line in hu_stock:

            # print(read_line)

            if data in read_line[0:6]:  # 判断是否匹配输入号码

                tag = True

                print(read_line)

                ret = read_line

                continue

            if read_line.startswith('60') and tag:  # 判断结束条件

                break

            # if tag:

            #    print(read_line)

            #    ret.append(read_line)

            # elif data not in hu_stock:  # 查询无记录

            #    print('查询无记录。')

    return ret

def add(data):  # 修改功能函数

    res = fetch(data)

    def change(res):

        change_w = input('请输入替换内容：') #替换内容

        change_w = change_w + '\n' #添加内容的换行

        # print(change_w)

        with open('沪市股票', 'r') as hu_stock, open('沪市股票-new', 'w')as new_stock:

            for read_line in hu_stock:

                tag = False

                if read_line == res: #判断查询内容

                    tag = True

                    new_stock.write(change_w) #把替换内容写入

                if read_line.startswith('60') and tag:

                    tag = False

                    continue

                if not tag:

                    new_stock.write(read_line) #正常逐行写入文件

        os.rename('沪市股票','沪市股票.bak') #把原文件改为备份

        os.rename('沪市股票-new','沪市股票') #把新文件改成原文件

        # os.remove('沪市股票.bak')

        return new_stock

    return change(res)