Refactor main script to utilize DiscreteMaterialPlanExtractor for data extraction and streamline order processing

main.py

@@ -1,196 +1,28 @@
 """
-主程序 - 批量下载离散备料计划维护数据
+主程序 - 使用离散备料计划维护数据提取工具
 """
-from playwright.sync_api import sync_playwright
-from login import login
-import re
 import os
-import pandas as pd
+from utils.离散备料计划维护数据提取 import DiscreteMaterialPlanExtractor
-from time import sleep
-from utils.excel_converter import ExcelConverter
-
-
-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} 批查询完成，等待加载结果...")
-
-    # 等待加载完成
-    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 convert_and_merge_files(file_paths, output_path):
-    """使用 ExcelConverter 转换并合并所有文件"""
-    converter = ExcelConverter(verbose=False)  # 不打印详细日志
-    all_dataframes = []
-
-    for i, file_path in enumerate(file_paths, 1):
-        print(f"转换第 {i} 个文件: {file_path}")
-        df = converter.convert(file_path, output_file=None)  # 只转换，不保存
-        all_dataframes.append(df)
-        print(f"  提取到 {len(df)} 条记录")
-
-    if all_dataframes:
-        print(f"\n合并 {len(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:
+    # 创建提取器
-        # 调用登录模块
+    extractor = DiscreteMaterialPlanExtractor(
-        browser, context, page, main_frame = login(
-            playwright=playwright,
         username="BLDpengqiangqiang",
         password="Cqbld123456.",
         headless=False,
-            ignore_https_errors=True
+        verbose=True
     )
 
-        # 登录成功后可以进行后续操作
+    # 设置文件路径
-        # 点击打开"菜单"
-        main_frame.locator("i").first.click()
-
-        # 点击打开"离散备料计划维护"
-        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()
-
-        # 填充并验证，如果失败则重试
-        max_retries = 3
-        expected_value = "5000"
-        for attempt in range(max_retries):
-            inner_frame.locator("#rc_select_0").fill(expected_value)
-            inner_frame.locator("#rc_select_0").press("Enter")
-            # 检查填充是否成功
-            actual_value = inner_frame.locator("#rc_select_0").input_value()
-            if actual_value == expected_value:
-                print(f"文本框填充成功: {expected_value}")
-                break
-            else:
-                print(f"第 {attempt + 1} 次填充失败，实际值: {actual_value}，重试...")
-                if attempt == max_retries - 1:
-                    print(f"警告: {max_retries} 次尝试后仍未成功填充，继续执行...")
-
-        # 读取订单号文件
     order_id_file = os.path.join(os.path.dirname(__file__), "orderID.txt")
-        order_ids = read_order_ids(order_id_file)
+    output_file = r"D:/python/playwrite/data/离散备料计划维护_合并.xlsx"
-        print(f"共读取到 {len(order_ids)} 个订单号")
 
-        # 确保 data 目录存在
+    # 执行提取
-        data_dir = "D:/python/playwrite/data"
+    extractor.extract(order_id_file, output_file)
-        os.makedirs(data_dir, exist_ok=True)
-
-        # 按批次下载
-        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=None)
-            downloaded_files.append(downloaded_file)
-
-        # 转换并合并文件
-        if downloaded_files:
-            print(f"\n=== 开始转换并合并 {len(downloaded_files)} 个文件 ===")
-            final_output = "D:/python/playwrite/data/离散备料计划维护_合并.xlsx"
-            convert_and_merge_files(downloaded_files, final_output)
-        else:
-            print("\n没有下载到任何文件")
-
-        print(f"\n=== 全部完成 ===")
-        print(f"最终文件: {final_output}")
 
     input("按回车退出...")
 
-        # 关闭浏览器
-        context.close()
-        browser.close()
-
 
 if __name__ == "__main__":
     main()

utils/__init__.py

@@ -2,5 +2,6 @@
 工具组件包
 """
 from .excel_converter import ExcelConverter
+from .离散备料计划维护数据提取 import DiscreteMaterialPlanExtractor
 
-__all__ = ['ExcelConverter']
+__all__ = ['ExcelConverter', 'DiscreteMaterialPlanExtractor']

utils/离散备料计划维护数据提取.py

@@ -0,0 +1,273 @@
+"""
+离散备料计划维护数据提取工具
+负责登录、批量下载、转换数据
+"""
+import os
+import pandas as pd
+from playwright.sync_api import sync_playwright
+from utils.excel_converter import ExcelConverter
+
+
+class DiscreteMaterialPlanExtractor:
+    """离散备料计划维护数据提取器"""
+
+    def __init__(self, username, password, headless=False, verbose=True):
+        """
+        初始化提取器
+
+        Args:
+            username: 登录用户名
+            password: 登录密码
+            headless: 是否无头模式运行
+            verbose: 是否打印详细日志
+        """
+        self.username = username
+        self.password = password
+        self.headless = headless
+        self.verbose = verbose
+        self.converter = ExcelConverter(verbose=verbose)
+
+    def _print(self, *args, **kwargs):
+        """打印日志（如果 verbose=True）"""
+        if self.verbose:
+            print(*args, **kwargs)
+
+    def read_order_ids(self, 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(self, 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(self, inner_frame, order_ids, batch_index, page1, debug_mode=False, debug_batch=None):
+        """下载一批订单号的数据"""
+        from playwright.sync_api import TimeoutError
+        import re
+        import time
+
+        # 清空文本框
+        textbox = inner_frame.get_by_role("textbox", name="来源生产订单号")
+        textbox.fill("")
+
+        # 填充订单号
+        textbox.fill(",".join(order_ids))
+
+        # 点击查询
+        inner_frame.locator(".search-component-searchBtn").click()
+        self._print(f"第 {batch_index + 1} 批查询完成，等待加载结果...")
+
+        # 等待加载完成
+        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
+        self._print(f"第 {batch_index + 1} 批加载完成，开始选择数据...")
+
+        # 调试模式：只在指定批次暂停
+        if debug_mode and (debug_batch is None or batch_index == debug_batch):
+            self._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)
+        self._print(f"第 {batch_index + 1} 批下载完成: {download_path}")
+
+        # 关闭输出对话框（如果有的话）
+        # try:
+        #     inner_frame.get_by_role("button", name="取消").click()
+        # except:
+        #     pass
+
+        # 等待页面恢复，准备下一次查询
+        time.sleep(1)
+
+        return download_path
+
+    def convert_and_merge_files(self, file_paths, output_path):
+        """使用 ExcelConverter 转换并合并所有文件"""
+        # 确保输出文件路径是正确的格式
+        output_path = os.path.normpath(output_path)
+        output_dir = os.path.dirname(output_path)
+        output_filename = os.path.basename(output_path)
+
+        self._print(f"输出文件路径: {output_path}")
+        self._print(f"输出目录: {output_dir}")
+        self._print(f"输出文件名: {output_filename}")
+
+        # 确保输出文件的父目录存在
+        if output_dir and not os.path.exists(output_dir):
+            self._print(f"创建输出目录: {output_dir}")
+            os.makedirs(output_dir)
+
+        all_dataframes = []
+
+        for i, file_path in enumerate(file_paths, 1):
+            self._print(f"转换第 {i} 个文件: {file_path}")
+            df = self.converter.convert(file_path, output_file=None)  # 只转换，不保存
+            all_dataframes.append(df)
+            self._print(f"  提取到 {len(df)} 条记录")
+
+        if all_dataframes:
+            self._print(f"\n合并 {len(all_dataframes)} 个文件的数据...")
+            merged_df = pd.concat(all_dataframes, ignore_index=True)
+            merged_df.to_excel(output_path, index=False)
+            self._print(f"合并完成: {output_path}, 总共 {len(merged_df)} 条记录")
+
+            # 删除临时文件
+            for file_path in file_paths:
+                os.remove(file_path)
+                self._print(f"已删除临时文件: {file_path}")
+
+            return output_path
+        return None
+
+    def setup_query_interface(self, inner_frame):
+        """设置查询界面"""
+        import re
+
+        # 点击图标按钮打开查询界面
+        inner_frame.locator(".search-name-wrapper > .iconfont").click()
+        inner_frame.get_by_text("订单号查询").click()
+        inner_frame.get_by_role("tab", name="全部").click()
+
+        # 填充并验证，如果失败则重试
+        max_retries = 3
+        expected_value = "5000"
+        for attempt in range(max_retries):
+            inner_frame.locator("#rc_select_0").fill(expected_value)
+            inner_frame.locator("#rc_select_0").press("Enter")
+            # 检查填充是否成功
+            actual_value = inner_frame.locator("#rc_select_0").input_value()
+            if actual_value == expected_value:
+                self._print(f"文本框填充成功: {expected_value}")
+                break
+            else:
+                self._print(f"第 {attempt + 1} 次填充失败，实际值: {actual_value}，重试...")
+                if attempt == max_retries - 1:
+                    self._print(f"警告: {max_retries} 次尝试后仍未成功填充，继续执行...")
+
+    def extract(self, order_id_file, data_dir="D:/python/playwrite/data",
+               output_file="D:/python/playwrite/data/离散备料计划维护_合并.xlsx",
+               debug_mode=False, debug_batch=None):
+        """
+        执行完整的数据提取流程
+
+        Args:
+            order_id_file: 订单号文件路径
+            data_dir: 数据保存目录
+            output_file: 最终输出文件路径
+            debug_mode: 是否启用调试模式
+            debug_batch: 调试批次号
+
+        Returns:
+            输出文件路径
+        """
+        from login import login
+
+        with sync_playwright() as playwright:
+            # 调用登录模块
+            browser, context, page, main_frame = login(
+                playwright=playwright,
+                username=self.username,
+                password=self.password,
+                headless=self.headless,
+                ignore_https_errors=True
+            )
+
+            self._print("=" * 80)
+            self._print("开始执行离散备料计划维护数据提取")
+            self._print("=" * 80)
+
+            # 登录成功后可以进行后续操作
+            # 点击打开"菜单"
+            main_frame.locator("i").first.click()
+
+            # 点击打开"离散备料计划维护"
+            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
+
+            # 设置查询界面
+            self.setup_query_interface(inner_frame)
+
+            # 读取订单号文件
+            order_ids = self.read_order_ids(order_id_file)
+            self._print(f"共读取到 {len(order_ids)} 个订单号")
+
+            # 按批次下载
+            downloaded_files = []
+            for batch_index, order_ids_batch in enumerate(self.group_order_ids(order_ids, 100)):
+                self._print(f"\n=== 开始处理第 {batch_index + 1} 批，共 {len(order_ids_batch)} 个订单号 ===")
+                downloaded_file = self.download_batch(
+                    inner_frame, order_ids_batch, batch_index, page1,
+                    debug_mode=debug_mode, debug_batch=debug_batch
+                )
+                downloaded_files.append(downloaded_file)
+
+
+            # 转换并合并文件
+            if downloaded_files:
+                self._print(f"\n=== 开始转换并合并 {len(downloaded_files)} 个文件 ===")
+                self.convert_and_merge_files(downloaded_files, output_file)
+            else:
+                self._print("\n没有下载到任何文件")
+
+            self._print(f"\n=== 全部完成 ===")
+            self._print(f"最终文件: {output_file}")
+
+            # 关闭浏览器
+            context.close()
+            browser.close()
+
+            return output_file
+
+
+def main():
+    """测试函数"""
+    extractor = DiscreteMaterialPlanExtractor(
+        username="BLDpengqiangqiang",
+        password="Cqbld123456.",
+        headless=False,
+        verbose=True
+    )
+
+    order_id_file = os.path.join(os.path.dirname(__file__), "orderID.txt")
+    output_file = "D:/python/playwrite/data/离散备料计划维护_合并.xlsx"
+
+    extractor.extract(order_id_file, output_file)
+
+    input("按回车退出...")
+
+
+if __name__ == "__main__":
+    main()