python协程学习

python 复制代码
import asyncio
import time
import csv
import queue
import aiosqlite
import time

conn = None
# 定义一个队列，用于传递数据
data_queue = queue.Queue()

# 启动写文件
# def callback():
#     await 
#     print(f"执行结果:{future.result()}")
async def get_connection():
    global conn
    if conn is None:
        conn = await aiosqlite.connect('lac_ci.db', check_same_thread=False)  # 在内存中创建一个临时的数据库
    return conn

async def query(row):
   
    lac = row[0]
    ci = row[1]
    
    # print("enter query")
    conn = await get_connection()
    cursor = await conn.execute("SELECT lon,lat,radius FROM LacCiT WHERE lac = ? AND ci = ?",(lac, ci))
    result = await cursor.fetchall()
    # print(result)
    # 获取查询结果的列名
        
    if(len(result)!=0): 
        result = result[0]   
        lat = result[0]
        lon = result[1]
        radius = result[2]
        # 将结果写入
        row.extend([str(lon),str(lat),str(radius)])
        # print(row)
    
    # 将处理后的行放入队列
    data_queue.put(row)
  
# 定义一个函数，用于写入数据到CSV文件
async def write_csv(filename):
    with open(filename, 'w', newline='') as file:
        writer = csv.writer(file)
        while True:
            try:
                # 从队列中获取处理后的行并写入文件
                row = data_queue.get(timeout=1)  # 设置超时以避免无限等待
                writer.writerow(row)
                data_queue.task_done()
            except queue.Empty:
                print("队列为空，说明数据已经全部处理完毕")
                break
    
async def main():
    
    start_time = time.time()  # 记录程序开始时间
    print(f"started at {time.strftime('%X')}")
    csv_file = 'lac_tt.csv'
    output_file = 'new_file.csv'
    
    
    file = open(csv_file, 'r', newline='') 
    reader = csv.reader(file)
    # row_num = reader.line_num
    # print(type(reader))
    tasks = []
    for row in reader:
        task = asyncio.create_task(query(row))
        tasks.append(task)
        
    
    # tasks.append(write_csv(output_file))
    await asyncio.gather(*tasks)
    # done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
    await write_csv(output_file)
    
    print(f"finished at {time.strftime('%X')}")
    
    end_time = time.time()  # 记录程序结束时间
    
    elapsed_time = end_time - start_time  # 计算运行时间
    print(f"程序运行时间：{elapsed_time:.2f} 秒")

    await conn.close()
    
coro = main() # 返回一个coroutine object

# 进入 event loop
asyncio.run(coro)