PDF关键字命中测试脚本

import PyPDF2import csvimport osfrom tqdm import tqdmfrom typing import Callabledef match_pdf(pdf_path: str, match_func: Callable[[str], tuple[bool, str]]) -> list[str]:    """    从指定的 PDF 文件中提取所有文本内容。    :param pdf_path: PDF 文件的路径    :param match_func: 匹配函数    :return: 包含 PDF 文件中所有文本的字符串    """    result = []    # 打开 PDF 文件    with open(pdf_path, 'rb') as file:        # 创建 PDF 阅读器对象        reader = PyPDF2.PdfReader(file)        # 遍历每一页并提取文本        for page in reader.pages:            match = match_func(page.extract_text())            if match[0]:                result.append(match[1])    return resultdef read_cities_from_csv(csv_path: str) -> list[str]:    """    从 CSV 文件中读取城市列表。    :param csv_path: CSV 文件的路径    :return: 包含城市列表的列表    """    with open(csv_path, 'r') as file:        reader = csv.reader(file)        return [row[0] for row in reader]def read_province_from_csv(csv_path: str) -> str:    """    从 CSV 文件中读取省份。    :param csv_path: CSV 文件的路径    :return: 省份    """    with open(csv_path, 'r') as file:        reader = csv.reader(file)        return [row[2] for row in reader]def match_keyword_in_pdf(pdf_path: str, keywords: list[str]) -> list[str]:    """    在 PDF 文件中匹配城市列表。    :param pdf_path: PDF 文件的路径    :param keywords: 关键词列表    :return: 包含匹配到的关键词的列表    """    def match_func(text: str) -> tuple[bool, str]:        for keyword in keywords:            if keyword in text:                return True, keyword        return False, ""    return match_pdf(pdf_path, match_func)if __name__ == "__main__":    input_dir = "./input"    cities_csv_path = "./cities.csv"    cities = read_cities_from_csv(cities_csv_path)    provinces = read_province_from_csv(cities_csv_path)    for file in tqdm(os.listdir(input_dir), desc="Processing files"):        # 遍历输入目录中的所有 PDF 文件        if file.endswith(".pdf"):            pdf_path = os.path.join(input_dir, file)            matched_cities = match_keyword_in_pdf(pdf_path, cities)            print(f"File {file} has {len(matched_cities)} cities, {matched_cities}")            matched_provinces = match_keyword_in_pdf(pdf_path, provinces)            print(f"File {file} has {len(matched_provinces)} provinces, {matched_provinces}")