Refactor main script to enhance batch downloading of order data and implement file merging functionality

main.py

@@ -1,11 +1,107 @@
 """
-主程序 - 演示如何使用登录模块
+主程序 - 批量下载离散备料计划维护数据
 """
 from playwright.sync_api import sync_playwright
 from login import login
-from time import sleep
 import re
+import os
+import pandas as pd
+from time import sleep
+
+def read_order_ids(file_path):
+    """读取订单号文件"""
+    with open(file_path, 'r', encoding='utf-8') as f:
+        # 去除空白行和空格
+        order_ids = [line.strip() for line in f if line.strip()]
+    return order_ids
+
+def group_order_ids(order_ids, group_size=100):
+    """将订单号分组"""
+    for i in range(0, len(order_ids), group_size):
+        yield order_ids[i:i + group_size]
+
+def download_batch(inner_frame, order_ids, batch_index, page1, debug_mode=False, debug_batch=None):
+    """下载一批订单号的数据"""
+    # 清空文本框
+    textbox = inner_frame.get_by_role("textbox", name="来源生产订单号")
+    textbox.fill("")
+
+    # 填充订单号
+    textbox.fill(",".join(order_ids))
+
+    # 点击查询
+    inner_frame.locator(".search-component-searchBtn").click()
+    print(f"第 {batch_index + 1} 批查询完成，等待加载结果...")
+    #sleep(3)
+    # 等待加载完成
+    loading_locator = inner_frame.locator("div").filter(has_text="加载中").nth(1)
+    try:
+        loading_locator.wait_for(state="visible", timeout=3000)
+        loading_locator.wait_for(state="hidden", timeout=0)  # 无限等待，直到消失
+    except TimeoutError:
+        # 加载很快完成，或者没有出现加载提示
+        pass
+    print(f"第 {batch_index + 1} 批加载完成，开始选择数据...")
+
+    # 调试模式：只在指定批次暂停
+    if debug_mode and (debug_batch is None or batch_index == debug_batch):
+        print(f"=== 调试暂停：第 {batch_index + 1} 批 ===")
+        page1.pause()
+
+    # 选择所有数据
+    inner_frame.get_by_role("row", name="序号").get_by_label("").click()
+
+    # 点击输出
+    inner_frame.get_by_role("button", name="更多").hover()
+    inner_frame.get_by_text("输出", exact=True).click()
+
+    # 设置行数阈值
+    input_box = inner_frame.locator("div").filter(has_text=re.compile(r"^行数阈值$")).locator("input[type='text']")
+    input_box.fill("300000")
+
+    # 下载文件
+    download_path = f"D:/python/playwrite/data/temp_batch_{batch_index + 1}.xlsx"
+    with page1.expect_download() as download_info:
+        inner_frame.get_by_role("button", name="确定(Y)").click()
+
+    download = download_info.value
+    download.save_as(download_path)
+    print(f"第 {batch_index + 1} 批下载完成: {download_path}")
+
+    # 关闭输出对话框（如果有的话）
+    try:
+        inner_frame.get_by_role("button", name="取消").click()
+    except:
+        pass
+
+    # 等待页面恢复，准备下一次查询
     import time
+    time.sleep(1)
+
+    return download_path
+
+def merge_excel_files(file_paths, output_path):
+    """合并多个 Excel 文件"""
+    all_dataframes = []
+
+    for file_path in file_paths:
+        df = pd.read_excel(file_path)
+        all_dataframes.append(df)
+        print(f"已读取: {file_path}, 共 {len(df)} 行")
+
+    if all_dataframes:
+        merged_df = pd.concat(all_dataframes, ignore_index=True)
+        merged_df.to_excel(output_path, index=False)
+        print(f"合并完成: {output_path}, 总共 {len(merged_df)} 行")
+
+        # 删除临时文件
+        for file_path in file_paths:
+            os.remove(file_path)
+            print(f"已删除临时文件: {file_path}")
+
+        return output_path
+    return None
+
 def main():
     with sync_playwright() as playwright:
         # 调用登录模块
@@ -18,59 +114,59 @@ def main():
         )
 
         # 登录成功后可以进行后续操作
-        # 点击打开“菜单”
+        # 点击打开"菜单"
         main_frame.locator("i").first.click()
-        #page.pause()
+
-        # 点击打开“菜单”
+        # 点击打开"离散备料计划维护"
         with page.expect_popup() as page1_info:
             main_frame.get_by_title("离散备料计划维护", exact=True).first.click()
         page1 = page1_info.value
 
-        
+        # 获取 nested iframe
         main_frame = page1.locator("#forwardFrame").content_frame
         inner_frame_locator = main_frame.locator("#mainiframe")
         inner_frame_locator.wait_for(state="visible", timeout=15000)
         inner_frame = inner_frame_locator.content_frame
 
-        # 点击图标按钮
+        # 点击图标按钮打开查询界面
         inner_frame.locator(".search-name-wrapper > .iconfont").click()
         inner_frame.get_by_text("订单号查询").click()
         inner_frame.get_by_role("tab", name="全部").click()
+        
         inner_frame.locator("#rc_select_0").fill("5000")
-        #发送回车键
         inner_frame.locator("#rc_select_0").press("Enter")
 
-        # 输入订单号
+        # 读取订单号文件
-        textbox = inner_frame.get_by_role("textbox", name="来源生产订单号")
+        order_id_file = os.path.join(os.path.dirname(__file__), "orderID.txt")
-        textbox.fill("SC70202510110003,SC70202510110004")
+        order_ids = read_order_ids(order_id_file)
-        inner_frame.locator(".search-component-searchBtn").click()
+        print(f"共读取到 {len(order_ids)} 个订单号")
-        print("查询完成，等待加载结果...")
-        # 等待加载完成
-        loading_locator = inner_frame.locator("div").filter(has_text="加载中").nth(1)
-        try:
-            loading_locator.wait_for(state="visible", timeout=3000)
-            loading_locator.wait_for(state="hidden", timeout=0)  # 无限等待，直到消失
-        except TimeoutError:
-            # 加载很快完成，或者没有出现加载提示
-            pass
-        print("加载完成，开始选择数据...")
-        inner_frame.get_by_role("row", name="序号").get_by_label("").check()
 
-        inner_frame.get_by_role("button", name="更多").hover()
+        # 确保 data 目录存在
-        inner_frame.get_by_text("输出").click()
+        data_dir = "D:/python/playwrite/data"
-        input_box = inner_frame.locator("div").filter(has_text=re.compile(r"^行数阈值$")).locator("input[type='text']")
+        os.makedirs(data_dir, exist_ok=True)
-        input_box.fill("300000")
-        # 等待下载事件
-        with page1.expect_download() as download_info:
-            inner_frame.get_by_role("button", name="确定(Y)").click()
 
-        # 获取下载对象
+        # 按批次下载
-        download = download_info.value
+        downloaded_files = []
+        for batch_index, order_ids_batch in enumerate(group_order_ids(order_ids, 100)):
+            print(f"\n=== 开始处理第 {batch_index + 1} 批，共 {len(order_ids_batch)} 个订单号 ===")
+            downloaded_file = download_batch(inner_frame, order_ids_batch, batch_index, page1, debug_mode=False, debug_batch=1)
+            downloaded_files.append(downloaded_file)
 
-        # 指定保存路径和文件名
+        # 合并文件
-        download.save_as("D:/python/playwrite/data/导出文件.xlsx")
+        if len(downloaded_files) > 1:
+            print(f"\n=== 开始合并 {len(downloaded_files)} 个文件 ===")
+            final_output = "D:/python/playwrite/data/离散备料计划维护_合并.xlsx"
+            merge_excel_files(downloaded_files, final_output)
+        else:
+            print("\n只下载了一个文件，无需合并")
+            final_output = downloaded_files[0]
+            # 重命名为最终文件名
+            os.rename(final_output, "D:/python/playwrite/data/离散备料计划维护_合并.xlsx")
+            final_output = "D:/python/playwrite/data/离散备料计划维护_合并.xlsx"
+
+        print(f"\n=== 全部完成 ===")
+        print(f"最终文件: {final_output}")
 
-        page1.pause()
         input("按回车退出...")
 
         # 关闭浏览器