xlwings/vba_test_runner.py

"""
VBA 自动化测试与精确报错系统
通过代码编织技术实现 VBA 宏的精确行号报错定位
"""

import re
import sys
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
import xlwings as xw


@dataclass
class TestResult:
    """测试结果数据类"""
    procedure_name: str
    success: bool
    error_number: int = 0
    error_description: str = ""
    error_line: int = 0
    source_code: str = ""


class CodeWeaver:
    """代码编织器 - 为 VBA 代码注入行号标签和错误处理"""

    def __init__(self):
        self.line_counter = 10  # 行号标签从 10 开始，步长 10

    def parse_procedures(self, code: str) -> Dict[str, Dict]:
        """
        解析 VBA 代码中的所有 Sub/Function

        返回: {
            过程名: {
                'start_line': int,
                'end_line': int,
                'type': 'Sub' or 'Function'
            }
        }
        """
        procedures = {}
        lines = code.split('\n')

        # 正则匹配 Sub/Function 定义行
        proc_pattern = re.compile(
            r'^\s*(Public|Private|Friend)?\s*(Sub|Function)\s+(\w+)\s*\((.*?)\)(?:\s+As\s+(\w+))?',
            re.IGNORECASE | re.MULTILINE
        )

        end_pattern = re.compile(r'^\s*End\s+(Sub|Function)', re.IGNORECASE)

        current_proc = None
        proc_start_line = 0

        for i, line in enumerate(lines):
            # 检查是否是过程定义
            match = proc_pattern.match(line)
            if match:
                visibility, proc_type, proc_name = match.groups()[:3]
                current_proc = proc_name
                proc_start_line = i
                continue

            # 检查是否是过程结束
            if current_proc:
                end_match = end_pattern.match(line)
                if end_match:
                    procedures[current_proc] = {
                        'start_line': proc_start_line,
                        'end_line': i,
                        'type': proc_type
                    }
                    current_proc = None

        return procedures

    def weave_procedure(self, code: str, proc_name: str) -> Tuple[str, Dict[int, str]]:
        """
        为指定过程注入行号和错误处理

        返回: (编织后的代码, {行号: 源代码})
        """
        procedures = self.parse_procedures(code)

        if proc_name not in procedures:
            raise ValueError(f"Procedure '{proc_name}' not found in code")

        proc_info = procedures[proc_name]
        lines = code.split('\n')

        # 提取过程代码（包括定义行和结束行）
        proc_lines = lines[proc_info['start_line']:proc_info['end_line'] + 1]

        # 注入行号
        woven_lines, source_map = self._inject_line_numbers(proc_lines)

        # 注入错误处理
        woven_lines = self._inject_error_handler(woven_lines, proc_name)

        # 替换原代码中的过程
        result_lines = (
            lines[:proc_info['start_line']] +
            woven_lines +
            lines[proc_info['end_line'] + 1:]
        )

        return '\n'.join(result_lines), source_map

    def _inject_line_numbers(self, lines: List[str]) -> Tuple[List[str], Dict[int, str]]:
        """
        注入行号标签（10, 20, 30...）
        跳过：
        - Dim 声明
        - 注释行
        - 空行
        - 现有的 On Error 语句
        - 过程定义和结束行
        """
        result_lines = []
        source_map = {}
        self.line_counter = 10

        # 检查是否是声明行（Dim, Private, Public 等）
        is_declaration = re.compile(r'^\s*(Dim|Private|Public|Friend|Const|Type|Declare)\b', re.IGNORECASE)
        is_comment = re.compile(r'^\s*\'')
        is_empty = re.compile(r'^\s*$')
        is_on_error = re.compile(r'^\s*On\s+Error', re.IGNORECASE)
        is_proc_def = re.compile(r'^\s*(Public|Private|Friend)?\s*(Sub|Function)', re.IGNORECASE)
        is_proc_end = re.compile(r'^\s*End\s+(Sub|Function)', re.IGNORECASE)
        is_exit_stmt = re.compile(r'^\s*Exit\s+(Sub|Function)', re.IGNORECASE)

        in_declaration_block = True

        for line in lines:
            # 检查是否声明块结束（遇到第一个可执行语句）
            if in_declaration_block and not is_declaration.match(line) and not is_comment.match(line) and not is_empty.match(line):
                if not is_proc_def.match(line) and not is_proc_end.match(line):
                    in_declaration_block = False

            # 判断是否需要注入行号
            should_inject = (
                not in_declaration_block and
                not is_comment.match(line) and
                not is_empty.match(line) and
                not is_on_error.match(line) and
                not is_proc_def.match(line) and
                not is_proc_end.match(line) and
                not is_exit_stmt.match(line)
            )

            if should_inject:
                # 保存源代码映射
                source_map[self.line_counter] = line.strip()
                # 注入行号标签
                result_lines.append(f"{self.line_counter}")
                result_lines.append(line)
                self.line_counter += 10
            else:
                result_lines.append(line)

        return result_lines, source_map

    def _inject_error_handler(self, lines: List[str], proc_name: str) -> List[str]:
        """
        在过程开头添加: On Error GoTo Auto_Err_Handler
        删除任何现有的 On Error 语句
        在过程结束前添加错误处理块
        """
        result_lines = []
        on_error_pattern = re.compile(r'^\s*On\s+Error.*$', re.IGNORECASE)
        proc_end_pattern = re.compile(r'^\s*End\s+(Sub|Function)', re.IGNORECASE)
        exit_pattern = re.compile(r'^\s*(Exit\s+(Sub|Function))', re.IGNORECASE)

        # 确定过程类型
        proc_type = "Sub"
        for line in lines:
            if re.match(r'^\s*(Public|Private|Friend)?\s*Function', line, re.IGNORECASE):
                proc_type = "Function"
                break

        # 在过程定义后插入 On Error 语句
        inserted_error_handler = False
        error_handler_inserted = False

        for i, line in enumerate(lines):
            # 跳过现有的 On Error 语句
            if on_error_pattern.match(line):
                continue

            # 检查是否是过程结束
            is_end = proc_end_pattern.match(line)

            if is_end and not error_handler_inserted:
                # 在 End Sub/Function 之前插入错误处理块
                result_lines.append(f"")
                result_lines.append(f"    Call TestLogger.LogSuccess()")
                if proc_type == "Sub":
                    result_lines.append(f"    Exit Sub")
                else:
                    result_lines.append(f"    Exit Function")
                result_lines.append(f"")
                result_lines.append(f"Auto_Err_Handler_{proc_name}:")
                result_lines.append(f"    Call TestLogger.LogError(\"{proc_name}\", Err.Number, Err.Description, Erl)")
                result_lines.append(line)  # 添加 End Sub/Function
                error_handler_inserted = True
            elif is_end:
                # 跳过后续的 End 语句（不应该有）
                pass
            elif exit_pattern.match(line):
                # 跳过现有的 Exit 语句
                pass
            else:
                # 在第一个非声明、非空行后插入 On Error
                if not inserted_error_handler and not re.match(r'^\s*(\'|Dim|Private|Public|Const|$)', line, re.IGNORECASE):
                    if not re.match(r'^\s*(Public|Private|Friend)?\s*(Sub|Function)', line, re.IGNORECASE):
                        result_lines.append(f"    On Error GoTo Auto_Err_Handler_{proc_name}")
                        inserted_error_handler = True

                result_lines.append(line)

        # 如果没有找到 End，追加到末尾
        if not error_handler_inserted:
            result_lines.append(f"")
            result_lines.append(f"    Call TestLogger.LogSuccess()")
            if proc_type == "Sub":
                result_lines.append(f"    Exit Sub")
            else:
                result_lines.append(f"    Exit Function")
            result_lines.append(f"")
            result_lines.append(f"Auto_Err_Handler_{proc_name}:")
            result_lines.append(f"    Call TestLogger.LogError(\"{proc_name}\", Err.Number, Err.Description, Erl)")

        return result_lines


class LoggerInjector:
    """日志模块注入器 - 向 Excel 项目注入 _TestLogger 辅助模块"""

    LOGGER_MODULE_CODE = """Option Explicit

' Module-level variables to store test result
Private m_TestStatus As String
Private m_ErrProcedure As String
Private m_ErrNumber As Long
Private m_ErrDescription As String
Private m_ErrLine As Long

Sub LogError(ByVal procName As String, ByVal errNum As Long, _
             ByVal errDesc As String, ByVal errLine As Long)
    m_TestStatus = "ERROR"
    m_ErrProcedure = procName
    m_ErrNumber = errNum
    m_ErrDescription = errDesc
    m_ErrLine = errLine
End Sub

Sub LogSuccess()
    m_TestStatus = "SUCCESS"
    m_ErrProcedure = ""
    m_ErrNumber = 0
    m_ErrDescription = ""
    m_ErrLine = 0
End Sub

Function GetResult() As String
    If m_TestStatus = "SUCCESS" Then
        GetResult = "SUCCESS"
    Else
        GetResult = "ERROR|" & m_ErrProcedure & "|" & m_ErrLine & "|" & _
                    m_ErrNumber & "|" & m_ErrDescription
    End If
End Function
"""

    def inject_or_replace(self, wb, module_name: str = "TestLogger") -> None:
        """
        如果模块存在则删除，然后注入新的 Logger 模块
        """
        # 获取 VBA 项目
        vba_project = wb.api.VBProject

        # 检查模块是否存在，存在则删除
        for component in vba_project.VBComponents:
            if component.Name == module_name:
                vba_project.VBComponents.Remove(component)
                break

        # 注入新模块
        new_module = vba_project.VBComponents.Add(1)  # 1 = vbext_ct_StdModule
        new_module.Name = module_name
        new_module.CodeModule.AddFromString(self.LOGGER_MODULE_CODE)


class TestRunner:
    """测试执行器 - 管理 Excel 生命周期，执行测试，返回结果"""

    def __init__(self, file_path: str, visible: bool = True):
        """
        初始化测试执行器

        Args:
            file_path: Excel 文件路径
            visible: 是否显示 Excel 窗口
        """
        self.file_path = file_path
        self.visible = visible
        self.app = None
        self.wb = None
        self.code_weaver = CodeWeaver()
        self.logger_injector = LoggerInjector()
        self.source_map = {}  # 行号 -> 源代码映射

    def run_test(self, module_name: str, proc_name: str) -> TestResult:
        """
        执行测试的完整流程:
        1. 打开 Excel 文件
        2. 读取原始 VBA 代码
        3. 使用 CodeWeaver 编织代码
        4. 注入/替换 Logger 模块
        5. 热替换目标模块代码
        6. 执行宏
        7. 获取结果
        8. 关闭工作簿（不保存）
        """
        try:
            # 1. 打开 Excel
            self.app = xw.App(visible=self.visible)
            self.wb = self.app.books.open(self.file_path)

            # 2. 读取原始 VBA 代码
            original_code = self._get_vba_code(module_name)

            # 3. 编织代码
            woven_code, self.source_map = self.code_weaver.weave_procedure(
                original_code, proc_name
            )

            # 4. 注入 Logger 模块
            self.logger_injector.inject_or_replace(self.wb)

            # 5. 热替换目标模块代码
            self._replace_module_code(module_name, woven_code)

            # 6. 执行宏
            self._execute_macro(f"{module_name}.{proc_name}")

            # 7. 获取结果
            result = self._get_test_result(proc_name)

            return result

        except Exception as e:
            return TestResult(
                procedure_name=proc_name,
                success=False,
                error_number=0,
                error_description=f"Test execution failed: {str(e)}",
                error_line=0,
                source_code=""
            )
        finally:
            # 8. 关闭工作簿（不保存）
            self._cleanup()

    def _get_vba_code(self, module_name: str) -> str:
        """从 Excel 读取指定模块的 VBA 代码"""
        vba_project = self.wb.api.VBProject

        for component in vba_project.VBComponents:
            if component.Name == module_name:
                code_module = component.CodeModule
                line_count = code_module.CountOfLines
                return code_module.Lines(1, line_count) if line_count > 0 else ""

        raise ValueError(f"Module '{module_name}' not found in VBA project")

    def _replace_module_code(self, module_name: str, new_code: str) -> None:
        """热替换模块代码"""
        vba_project = self.wb.api.VBProject

        for component in vba_project.VBComponents:
            if component.Name == module_name:
                code_module = component.CodeModule
                # 删除所有现有代码
                if code_module.CountOfLines > 0:
                    code_module.DeleteLines(1, code_module.CountOfLines)
                # 添加新代码
                code_module.AddFromString(new_code)
                return

        raise ValueError(f"Module '{module_name}' not found in VBA project")

    def _execute_macro(self, proc_name: str) -> None:
        """执行 VBA 宏"""
        try:
            self.app.api.Run(proc_name)
        except Exception as e:
            # 错误会被 VBA 的错误处理捕获，这里不需要处理
            pass

    def _get_test_result(self, proc_name: str) -> TestResult:
        """从 Logger 模块获取测试结果"""
        try:
            result_str = self.app.api.Run("TestLogger.GetResult")

            if result_str == "SUCCESS":
                return TestResult(
                    procedure_name=proc_name,
                    success=True
                )
            else:
                # 解析错误信息: ERROR|ProcedureName|Line|Number|Description
                parts = result_str.split('|')
                if len(parts) >= 5:
                    error_proc_name = parts[1]
                    error_line = int(parts[2])
                    error_number = int(parts[3])
                    error_desc = parts[4]

                    # 从 Source Map 获取源代码
                    source_code = self.source_map.get(error_line, "Source code not found")

                    return TestResult(
                        procedure_name=error_proc_name,
                        success=False,
                        error_number=error_number,
                        error_description=error_desc,
                        error_line=error_line,
                        source_code=source_code
                    )
        except Exception as e:
            return TestResult(
                procedure_name=proc_name,
                success=False,
                error_number=0,
                error_description=f"Failed to get test result: {str(e)}",
                error_line=0,
                source_code=""
            )

    def _cleanup(self):
        """清理资源，关闭 Excel"""
        try:
            if self.wb:
                self.wb.api.Close(False)
            if self.app:
                self.app.quit()
        except:
            pass


def print_test_result(result: TestResult, index: int = None, total: int = None):
    """打印单个测试结果"""
    prefix = f"[{index}/{total}] " if index and total else ""

    if result.success:
        print(f"{prefix}[PASS] {result.procedure_name} - Test Passed")
    else:
        print(f"{prefix}[FAIL] {result.procedure_name} - Test Failed")
        print(f"  Error Description: {result.error_description}")
        print(f"  Error Line: {result.error_line}")
        print(f"  Source Code: {result.source_code}")


def print_summary(results: List[TestResult]):
    """打印测试汇总"""
    passed = sum(1 for r in results if r.success)
    failed = sum(1 for r in results if not r.success)
    total = len(results)

    print("\n===== Test Summary =====")
    print(f"Passed: {passed}/{total}")
    print(f"Failed: {failed}/{total}")


def main():
    """主入口函数"""
    if len(sys.argv) < 3:
        print("Usage: python vba_test_runner.py <excel_file> <module_name> <procedure1> [procedure2] ...")
        print("Example: python vba_test_runner.py demo.xlsm Module1 TestErrorProcedure")
        sys.exit(1)

    excel_file = sys.argv[1]
    module_name = sys.argv[2]
    procedures = sys.argv[3:]

    print(f"===== VBA Batch Test Results =====\n")

    results = []

    for i, proc_name in enumerate(procedures, 1):
        runner = TestRunner(excel_file, visible=False)
        result = runner.run_test(module_name, proc_name)
        results.append(result)
        print_test_result(result, i, len(procedures))

    print_summary(results)


if __name__ == "__main__":
    main()