pdf_extractor.py 7.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246
  1. #!/usr/bin/env python3
  2. """
  3. PDF页面提取工具
  4. 从PDF文件中提取指定页面并保存为新PDF文件。
  5. 使用方法:
  6. python pdf_extractor.py input.pdf --pages "1-5,7,9-12" --output output.pdf
  7. python pdf_extractor.py input.pdf --pages "1-" --output output.pdf # 提取第1页到最后
  8. python pdf_extractor.py input.pdf --pages "-10" --output output.pdf # 提取前10页
  9. """
  10. import argparse
  11. import sys
  12. from pathlib import Path
  13. from typing import List
  14. import io
  15. try:
  16. import pypdfium2 as pdfium
  17. PDFIUM_AVAILABLE = True
  18. except ImportError:
  19. PDFIUM_AVAILABLE = False
  20. pdfium = None
  21. from loguru import logger
  22. from core.pdf_utils import PDFUtils
  23. def extract_pdf_pages(
  24. input_pdf_path: Path,
  25. page_indices: List[int],
  26. output_pdf_path: Path
  27. ) -> bool:
  28. """
  29. 从PDF中提取指定页面并保存为新PDF
  30. Args:
  31. input_pdf_path: 输入PDF文件路径
  32. page_indices: 要提取的页面索引列表(0-based)
  33. output_pdf_path: 输出PDF文件路径
  34. Returns:
  35. 是否成功
  36. """
  37. if not PDFIUM_AVAILABLE:
  38. logger.error("❌ pypdfium2 未安装,请先安装: pip install pypdfium2")
  39. return False
  40. if not input_pdf_path.exists():
  41. logger.error(f"❌ 输入文件不存在: {input_pdf_path}")
  42. return False
  43. if not input_pdf_path.suffix.lower() == '.pdf':
  44. logger.error(f"❌ 输入文件不是PDF格式: {input_pdf_path}")
  45. return False
  46. try:
  47. # 读取PDF文件
  48. with open(input_pdf_path, 'rb') as f:
  49. pdf_bytes = f.read()
  50. # 加载PDF文档
  51. pdf = pdfium.PdfDocument(pdf_bytes)
  52. total_pages = len(pdf)
  53. if total_pages == 0:
  54. logger.error("❌ PDF文件为空")
  55. pdf.close()
  56. return False
  57. # 验证页面索引
  58. valid_indices = []
  59. for idx in sorted(set(page_indices)): # 去重并排序
  60. if 0 <= idx < total_pages:
  61. valid_indices.append(idx)
  62. else:
  63. logger.warning(f"⚠️ 页面索引 {idx + 1} 超出范围(总页数: {total_pages}),已跳过")
  64. if not valid_indices:
  65. logger.error("❌ 没有有效的页面可提取")
  66. pdf.close()
  67. return False
  68. # 创建新PDF文档
  69. output_pdf = pdfium.PdfDocument.new()
  70. # 导入指定页面
  71. success_count = 0
  72. for page_idx in valid_indices:
  73. try:
  74. output_pdf.import_pages(pdf, pages=[page_idx])
  75. success_count += 1
  76. except Exception as e:
  77. logger.warning(f"⚠️ 导入第 {page_idx + 1} 页失败: {e},已跳过")
  78. continue
  79. if success_count == 0:
  80. logger.error("❌ 没有成功导入任何页面")
  81. pdf.close()
  82. output_pdf.close()
  83. return False
  84. # 保存到文件
  85. output_pdf_path.parent.mkdir(parents=True, exist_ok=True)
  86. # 保存到内存缓冲区
  87. output_buffer = io.BytesIO()
  88. output_pdf.save(output_buffer)
  89. output_bytes = output_buffer.getvalue()
  90. # 写入文件
  91. with open(output_pdf_path, 'wb') as f:
  92. f.write(output_bytes)
  93. # 清理资源
  94. pdf.close()
  95. output_pdf.close()
  96. logger.info(f"✅ 成功提取 {success_count} 页到: {output_pdf_path}")
  97. logger.info(f" 提取的页面: {', '.join([str(idx + 1) for idx in valid_indices])}")
  98. return True
  99. except Exception as e:
  100. logger.error(f"❌ 提取PDF页面时出错: {e}")
  101. import traceback
  102. logger.debug(traceback.format_exc())
  103. return False
  104. def main():
  105. """命令行入口"""
  106. parser = argparse.ArgumentParser(
  107. description='从PDF文件中提取指定页面并保存为新PDF文件',
  108. formatter_class=argparse.RawDescriptionHelpFormatter,
  109. epilog="""
  110. 示例:
  111. # 提取第1-5页和第7页
  112. python pdf_extractor.py --input input.pdf --pages "1-5,7" --output output.pdf
  113. # 提取第1页到最后
  114. python pdf_extractor.py --input input.pdf --pages "1-" --output output.pdf
  115. # 提取前10页
  116. python pdf_extractor.py --input input.pdf --pages "-10" --output output.pdf
  117. # 提取单页
  118. python pdf_extractor.py --input input.pdf --pages "3" --output output.pdf
  119. """
  120. )
  121. parser.add_argument(
  122. '--input', '-i',
  123. type=str,
  124. help='输入PDF文件路径'
  125. )
  126. parser.add_argument(
  127. '--pages', '-p',
  128. type=str,
  129. required=True,
  130. help='要提取的页面范围,支持格式: "1-5,7,9-12", "1-", "-10", "3"'
  131. )
  132. parser.add_argument(
  133. '--output', '-o',
  134. type=str,
  135. required=True,
  136. help='输出PDF文件路径'
  137. )
  138. args = parser.parse_args()
  139. # 解析输入路径
  140. input_path = Path(args.input).resolve()
  141. # 解析输出路径
  142. output_path = Path(args.output).resolve()
  143. # 如果输出路径是目录,自动生成文件名
  144. if output_path.is_dir() or not output_path.suffix.lower() == '.pdf':
  145. output_path = output_path / f"{input_path.stem}_extracted.pdf"
  146. # 检查 pypdfium2 是否可用
  147. if not PDFIUM_AVAILABLE:
  148. logger.error("❌ pypdfium2 未安装,请先安装: pip install pypdfium2")
  149. sys.exit(1)
  150. # 先加载PDF获取总页数(用于验证页面范围)
  151. try:
  152. with open(input_path, 'rb') as f:
  153. pdf_bytes = f.read()
  154. pdf = pdfium.PdfDocument(pdf_bytes)
  155. total_pages = len(pdf)
  156. pdf.close()
  157. except Exception as e:
  158. logger.error(f"❌ 无法读取PDF文件: {e}")
  159. sys.exit(1)
  160. # 解析页面范围
  161. page_set = PDFUtils.parse_page_range(args.pages, total_pages)
  162. page_indices = sorted(list(page_set))
  163. if not page_indices:
  164. logger.error(f"❌ 页面范围 '{args.pages}' 没有匹配到任何有效页面(总页数: {total_pages})")
  165. sys.exit(1)
  166. logger.info(f"📋 PDF总页数: {total_pages}")
  167. logger.info(f"📋 要提取的页面: {args.pages} → {len(page_indices)} 页")
  168. # 执行提取
  169. success = extract_pdf_pages(input_path, page_indices, output_path)
  170. if success:
  171. logger.info(f"✅ 提取完成!输出文件: {output_path}")
  172. sys.exit(0)
  173. else:
  174. logger.error("❌ 提取失败")
  175. sys.exit(1)
  176. if __name__ == '__main__':
  177. if len(sys.argv) == 1:
  178. # 没有命令行参数时,使用默认配置运行
  179. print("ℹ️ 未提供命令行参数,使用默认配置运行...")
  180. # 默认配置
  181. default_config = {
  182. "input": "/Users/zhch158/workspace/data/流水分析/施博深.pdf",
  183. "output": "/Users/zhch158/workspace/data/流水分析/施博深_extracted.pdf",
  184. # 页面范围(可选)
  185. "pages": "1-20", # 只处理前20页
  186. # "pages": "1-3,5,7-10", # 处理指定页面
  187. }
  188. # 构造参数
  189. sys.argv = [sys.argv[0]]
  190. for key, value in default_config.items():
  191. if isinstance(value, bool):
  192. if value:
  193. sys.argv.append(f"--{key}")
  194. else:
  195. sys.argv.extend([f"--{key}", str(value)])
  196. sys.exit(main())