tretool 0.2.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

tretool/transform/pdf.py

@@ -1,22 +1,70 @@
 import os
+import sys
+import logging
+import time
+import json
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Union, List, Optional
-from pdfminer.high_level import extract_text
-from pdf2image import convert_from_path
-from docx import Document
-import pandas as pd
-from tabula import read_pdf
-from PIL import Image
-import json
+from typing import Union, List, Optional, Dict, Callable
+from concurrent.futures import ThreadPoolExecutor
+from functools import wraps, lru_cache
+import argparse
+from datetime import datetime
 
+# 第三方库导入
+try:
+    from pdfminer.high_level import extract_text
+    from pdf2image import convert_from_path
+    from docx import Document
+    import pandas as pd
+    from tabula import read_pdf
+    from PIL import Image
+    from PyPDF2 import PdfReader
+except ImportError as e:
+    print(f"缺少依赖库: {e}")
+    sys.exit(1)
 
+# 日志配置
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler(),
+        logging.FileHandler('pdf_converter.log')
+    ]
+)
+
+# 自定义异常
 class ConversionError(Exception):
-    """自定义转换异常"""
+    """基础转换异常"""
+    pass
+
+class PDFPermissionError(ConversionError):
+    """PDF权限错误"""
+    pass
+
+class PDFCorruptedError(ConversionError):
+    """PDF文件损坏"""
     pass
 
+class UnsupportedFormatError(ConversionError):
+    """不支持的格式"""
+    pass
 
 class PDFConverter(ABC):
+    """
+    PDF文件转换器抽象基类
+    
+    提供PDF文件转换的基础功能，包括:
+    - 文件路径验证
+    - 输出路径处理
+    - 基本错误处理
+    
+    子类应实现:
+    - convert() 方法: 执行实际转换逻辑
+    - supported_formats() 类方法: 返回支持的格式列表
+    """
+    
     def __init__(self, pdf_path: Union[str, Path]):
         """
         初始化PDF转换器
@@ -25,19 +73,28 @@ class PDFConverter(ABC):
             pdf_path: PDF文件路径
         """
         self.pdf_path = Path(pdf_path)
+        self.logger = logging.getLogger(self.__class__.__name__)
+        self._progress_callback = None
+        
+        # 验证文件
         if not self.pdf_path.exists():
             raise FileNotFoundError(f"PDF文件不存在: {pdf_path}")
         if self.pdf_path.suffix.lower() != '.pdf':
             raise ValueError("输入文件必须是PDF格式")
+        
+        self.logger.info(f"初始化转换器，处理文件: {pdf_path}")
 
     @abstractmethod
-    def convert(self, output_path: Union[str, Path], **kwargs):
+    def convert(self, output_path: Union[str, Path], **kwargs) -> Union[str, List[str]]:
         """
         将PDF转换为目标格式
         
         参数:
             output_path: 输出文件路径
             **kwargs: 转换选项
+            
+        返回:
+            转换后的文件路径或路径列表
         """
         pass
 
@@ -47,6 +104,31 @@ class PDFConverter(ABC):
         """返回支持的格式列表"""
         return []
 
+    def set_progress_callback(self, callback: Callable[[int, int], None]):
+        """设置进度回调函数"""
+        self._progress_callback = callback
+        
+    def _update_progress(self, current: int, total: int):
+        """更新进度"""
+        if self._progress_callback:
+            self._progress_callback(current, total)
+
+    def get_metadata(self) -> Dict[str, str]:
+        """获取PDF元数据"""
+        try:
+            with open(self.pdf_path, 'rb') as f:
+                reader = PdfReader(f)
+                return {
+                    'title': reader.metadata.get('/Title', ''),
+                    'author': reader.metadata.get('/Author', ''),
+                    'pages': len(reader.pages),
+                    'created': reader.metadata.get('/CreationDate', ''),
+                    'modified': reader.metadata.get('/ModDate', '')
+                }
+        except Exception as e:
+            self.logger.warning(f"获取元数据失败: {str(e)}")
+            return {}
+
     def _prepare_output_path(self, output_path: Union[str, Path], 
                            default_extension: str) -> Path:
         """
@@ -73,9 +155,27 @@ class PDFConverter(ABC):
         
         return output_path
 
+# 重试装饰器
+def retry(max_attempts=3, delay=1, exceptions=(Exception,)):
+    def decorator(f):
+        @wraps(f)
+        def wrapper(*args, **kwargs):
+            last_error = None
+            for attempt in range(1, max_attempts+1):
+                try:
+                    return f(*args, **kwargs)
+                except exceptions as e:
+                    last_error = e
+                    if attempt < max_attempts:
+                        time.sleep(delay)
+            raise last_error
+        return wrapper
+    return decorator
 
+# 具体转换器实现
 class PDFToDocxConverter(PDFConverter):
-    def convert(self, output_path: Union[str, Path], **kwargs):
+    @retry(max_attempts=3, delay=0.5)
+    def convert(self, output_path: Union[str, Path], **kwargs) -> str:
         """
         将PDF转换为Word文档(.docx)
         
@@ -85,6 +185,9 @@ class PDFToDocxConverter(PDFConverter):
                 - start_page: 开始页(从1开始)
                 - end_page: 结束页
                 - preserve_formatting: 是否保留格式(True/False)
+                
+        返回:
+            转换后的文件路径
         """
         try:
             output_path = self._prepare_output_path(output_path, 'docx')
@@ -93,27 +196,35 @@ class PDFToDocxConverter(PDFConverter):
             end_page = kwargs.get('end_page', None)
             preserve = kwargs.get('preserve_formatting', False)
             
-            doc = Document()
-            text = extract_text(
-                str(self.pdf_path),
-                page_numbers=range(start_page-1, end_page) if end_page else None
-            )
-            
-            # 简单的段落处理
-            for paragraph in text.split('\n\n'):
-                if paragraph.strip():
-                    para = doc.add_paragraph()
-                    if preserve:
-                        # 这里可以添加更复杂的格式保留逻辑
-                        runs = paragraph.split('\n')
-                        for run in runs:
-                            if run.strip():
-                                para.add_run(run.strip() + ' ')
-                    else:
-                        para.add_run(paragraph.strip())
-            
-            doc.save(output_path)
-            return str(output_path)
+            # 尝试使用pdf2docx库（如果安装）
+            try:
+                from pdf2docx import Converter
+                cv = Converter(str(self.pdf_path))
+                cv.convert(str(output_path), start=start_page, end=end_page)
+                cv.close()
+                return str(output_path)
+            except ImportError:
+                self.logger.warning("pdf2docx未安装，使用基本文本转换")
+                # 回退到基本实现
+                doc = Document()
+                text = extract_text(
+                    str(self.pdf_path),
+                    page_numbers=range(start_page-1, end_page) if end_page else None
+                )
+                
+                for paragraph in text.split('\n\n'):
+                    if paragraph.strip():
+                        para = doc.add_paragraph()
+                        if preserve:
+                            runs = paragraph.split('\n')
+                            for run in runs:
+                                if run.strip():
+                                    para.add_run(run.strip() + ' ')
+                        else:
+                            para.add_run(paragraph.strip())
+                
+                doc.save(output_path)
+                return str(output_path)
         except Exception as e:
             raise ConversionError(f"转换为DOCX失败: {str(e)}")
 
@@ -121,9 +232,9 @@ class PDFToDocxConverter(PDFConverter):
     def supported_formats(cls) -> List[str]:
         return ['docx']
 
-
 class PDFToImageConverter(PDFConverter):
-    def convert(self, output_path: Union[str, Path], **kwargs):
+    @retry(max_attempts=3, delay=0.5)
+    def convert(self, output_path: Union[str, Path], **kwargs) -> Union[str, List[str]]:
         """
         将PDF转换为图像
         
@@ -134,6 +245,9 @@ class PDFToImageConverter(PDFConverter):
                 - fmt: 图像格式(png/jpg/tiff)
                 - merge: 是否合并所有页为一张长图(True/False)
                 - quality: 图像质量(1-100)
+                
+        返回:
+            单个文件路径或多个文件路径列表
         """
         try:
             dpi = kwargs.get('dpi', 200)
@@ -150,6 +264,8 @@ class PDFToImageConverter(PDFConverter):
                 fmt=fmt if fmt != 'jpg' else 'jpeg'
             )
             
+            self._update_progress(0, len(images))
+            
             if merge:
                 # 合并所有页为一张长图
                 output_path = self._prepare_output_path(output_path, fmt)
@@ -158,9 +274,10 @@ class PDFToImageConverter(PDFConverter):
                 
                 merged_image = Image.new('RGB', (max_width, total_height))
                 y_offset = 0
-                for img in images:
+                for i, img in enumerate(images):
                     merged_image.paste(img, (0, y_offset))
                     y_offset += img.height
+                    self._update_progress(i+1, len(images))
                 
                 merged_image.save(output_path, quality=quality)
                 return str(output_path)
@@ -170,25 +287,43 @@ class PDFToImageConverter(PDFConverter):
                 if len(images) == 1:
                     output_path = self._prepare_output_path(output_path, fmt)
                     images[0].save(output_path, quality=quality)
+                    self._update_progress(1, 1)
                     return str(output_path)
                 else:
                     output_path.mkdir(parents=True, exist_ok=True)
                     output_files = []
-                    for i, image in enumerate(images):
-                        page_path = output_path / f"page_{i+1}.{fmt}"
-                        image.save(page_path, quality=quality)
-                        output_files.append(str(page_path))
+                    
+                    # 多页并行处理
+                    if len(images) > 5:
+                        with ThreadPoolExecutor() as executor:
+                            results = list(executor.map(
+                                lambda x: self._save_image(x[1], output_path/f"page_{x[0]+1}.{fmt}", quality),
+                                enumerate(images)
+                            ))
+                            output_files.extend(results)
+                    else:
+                        for i, image in enumerate(images):
+                            page_path = output_path / f"page_{i+1}.{fmt}"
+                            image.save(page_path, quality=quality)
+                            output_files.append(str(page_path))
+                            self._update_progress(i+1, len(images))
+                    
                     return output_files
         except Exception as e:
             raise ConversionError(f"转换为图像失败: {str(e)}")
 
+    def _save_image(self, image: Image.Image, path: Path, quality: int) -> str:
+        """保存单个图像"""
+        image.save(path, quality=quality)
+        return str(path)
+
     @classmethod
     def supported_formats(cls) -> List[str]:
         return ['png', 'jpg', 'jpeg', 'tiff']
 
-
 class PDFToTextConverter(PDFConverter):
-    def convert(self, output_path: Union[str, Path], **kwargs):
+    @retry(max_attempts=3, delay=0.5)
+    def convert(self, output_path: Union[str, Path], **kwargs) -> str:
         """
         将PDF转换为纯文本
         
@@ -198,6 +333,9 @@ class PDFToTextConverter(PDFConverter):
                 - start_page: 开始页(从1开始)
                 - end_page: 结束页
                 - encoding: 文本编码(默认utf-8)
+                
+        返回:
+            转换后的文件路径
         """
         try:
             output_path = self._prepare_output_path(output_path, 'txt')
@@ -222,9 +360,9 @@ class PDFToTextConverter(PDFConverter):
     def supported_formats(cls) -> List[str]:
         return ['txt']
 
-
 class PDFToCSVConverter(PDFConverter):
-    def convert(self, output_path: Union[str, Path], **kwargs):
+    @retry(max_attempts=3, delay=0.5)
+    def convert(self, output_path: Union[str, Path], **kwargs) -> Union[str, List[str]]:
         """
         提取PDF中的表格为CSV
         
@@ -234,6 +372,9 @@ class PDFToCSVConverter(PDFConverter):
                 - pages: 要提取的页码('all'或数字或列表)
                 - multiple_tables: 如何处理多个表格(separate/merge)
                 - encoding: CSV文件编码(默认utf-8)
+                
+        返回:
+            单个CSV文件路径或多个CSV文件路径列表
         """
         try:
             pages = kwargs.get('pages', 'all')
@@ -245,11 +386,14 @@ class PDFToCSVConverter(PDFConverter):
             if not dfs:
                 raise ConversionError("未找到表格数据")
                 
+            self._update_progress(0, len(dfs))
+                
             if multiple_tables == 'merge':
                 # 合并所有表格
                 output_path = self._prepare_output_path(output_path, 'csv')
                 merged_df = pd.concat(dfs, ignore_index=True)
                 merged_df.to_csv(output_path, index=False, encoding=encoding)
+                self._update_progress(1, 1)
                 return str(output_path)
             else:
                 # 每个表格保存为单独CSV
@@ -257,6 +401,7 @@ class PDFToCSVConverter(PDFConverter):
                 if len(dfs) == 1:
                     output_path = self._prepare_output_path(output_path, 'csv')
                     dfs[0].to_csv(output_path, index=False, encoding=encoding)
+                    self._update_progress(1, 1)
                     return str(output_path)
                 else:
                     output_path.mkdir(parents=True, exist_ok=True)
@@ -265,6 +410,7 @@ class PDFToCSVConverter(PDFConverter):
                         table_path = output_path / f"table_{i+1}.csv"
                         df.to_csv(table_path, index=False, encoding=encoding)
                         output_files.append(str(table_path))
+                        self._update_progress(i+1, len(dfs))
                     return output_files
         except Exception as e:
             raise ConversionError(f"提取表格失败: {str(e)}")
@@ -273,9 +419,9 @@ class PDFToCSVConverter(PDFConverter):
     def supported_formats(cls) -> List[str]:
         return ['csv']
 
-
 class PDFToHTMLConverter(PDFConverter):
-    def convert(self, output_path: Union[str, Path], **kwargs):
+    @retry(max_attempts=3, delay=0.5)
+    def convert(self, output_path: Union[str, Path], **kwargs) -> str:
         """
         将PDF转换为HTML
         
@@ -284,6 +430,9 @@ class PDFToHTMLConverter(PDFConverter):
             **kwargs:
                 - css: 自定义CSS样式
                 - images: 是否嵌入图像(True/False)
+                
+        返回:
+            转换后的HTML文件路径
         """
         try:
             output_path = self._prepare_output_path(output_path, 'html')
@@ -298,7 +447,7 @@ class PDFToHTMLConverter(PDFConverter):
     <meta charset="UTF-8">
     <title>{self.pdf_path.stem}</title>
     <style>
-        {kwargs.get('css', 'body { font-family: Arial; margin: 20px; }')}
+        {kwargs.get('css', 'body {{ font-family: Arial; margin: 20px; }}')}
     </style>
 </head>
 <body>
@@ -320,9 +469,21 @@ class PDFToHTMLConverter(PDFConverter):
     def supported_formats(cls) -> List[str]:
         return ['html']
 
-
+# 工厂类
 class PDFConverterFactory:
+    _format_map = {
+        'docx': (PDFToDocxConverter, 'Microsoft Word文档'),
+        'txt': (PDFToTextConverter, '纯文本文件'),
+        'png': (PDFToImageConverter, 'PNG图像'),
+        'jpg': (PDFToImageConverter, 'JPEG图像'),
+        'jpeg': (PDFToImageConverter, 'JPEG图像'),
+        'tiff': (PDFToImageConverter, 'TIFF图像'),
+        'csv': (PDFToCSVConverter, 'CSV表格数据'),
+        'html': (PDFToHTMLConverter, 'HTML网页'),
+    }
+
     @staticmethod
+    @lru_cache(maxsize=32)
     def get_converter(target_format: str, pdf_path: str) -> PDFConverter:
         """
         获取指定格式的转换器
@@ -333,64 +494,81 @@ class PDFConverterFactory:
             
         返回:
             PDFConverter实例
+            
+        抛出:
+            UnsupportedFormatError: 当格式不支持时
         """
-        format_map = {
-            'docx': PDFToDocxConverter,
-            'txt': PDFToTextConverter,
-            'png': PDFToImageConverter,
-            'jpg': PDFToImageConverter,
-            'jpeg': PDFToImageConverter,
-            'tiff': PDFToImageConverter,
-            'csv': PDFToCSVConverter,
-            'html': PDFToHTMLConverter,
-        }
-        
         target_format = target_format.lower()
-        if target_format not in format_map:
-            raise ValueError(f"不支持的格式: {target_format}")
+        if target_format not in PDFConverterFactory._format_map:
+            raise UnsupportedFormatError(f"不支持的格式: {target_format}")
         
-        return format_map[target_format](pdf_path)
+        return PDFConverterFactory._format_map[target_format][0](pdf_path)
 
     @staticmethod
-    def get_supported_formats() -> dict:
-        """获取所有支持的格式"""
-        return {
-            'docx': 'Microsoft Word文档',
-            'txt': '纯文本文件',
-            'png': 'PNG图像',
-            'jpg': 'JPEG图像',
-            'jpeg': 'JPEG图像',
-            'tiff': 'TIFF图像',
-            'csv': 'CSV表格数据',
-            'html': 'HTML网页',
-        }
+    def get_supported_formats() -> Dict[str, str]:
+        """获取所有支持的格式及其描述"""
+        return {fmt: desc for fmt, (_, desc) in PDFConverterFactory._format_map.items()}
 
+# 命令行接口
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description='PDF转换工具 - 支持多种格式转换',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument('input', help='输入PDF文件路径')
+    parser.add_argument('output', help='输出文件路径或目录')
+    parser.add_argument('-f', '--format', required=True, 
+                       choices=PDFConverterFactory.get_supported_formats().keys(),
+                       help='目标格式')
+    parser.add_argument('--dpi', type=int, default=200, 
+                       help='图像DPI(仅图像转换)')
+    parser.add_argument('--start-page', type=int, default=1,
+                       help='起始页码(从1开始)')
+    parser.add_argument('--end-page', type=int, 
+                       help='结束页码(默认为最后一页)')
+    parser.add_argument('--preserve-formatting', action='store_true',
+                       help='保留格式(仅DOCX转换)')
+    parser.add_argument('--merge', action='store_true',
+                       help='合并多页为单个文件(图像/表格)')
+    parser.add_argument('--quality', type=int, default=90,
+                       help='输出质量(1-100, 仅图像)')
+    return parser.parse_args()
 
-# 使用示例
-if __name__ == "__main__":
+def main():
+    args = parse_args()
+    
     try:
-        # 示例1: PDF转Word
-        print("转换PDF到Word...")
-        docx_converter = PDFConverterFactory.get_converter('docx', 'example.pdf')
-        result = docx_converter.convert('output.docx', preserve_formatting=True)
-        print(f"转换成功: {result}")
+        # 获取转换器
+        converter = PDFConverterFactory.get_converter(args.format, args.input)
         
-        # 示例2: PDF转图像
-        print("\n转换PDF到图像...")
-        img_converter = PDFConverterFactory.get_converter('png', 'example.pdf')
-        result = img_converter.convert('output_images', dpi=300)
-        print(f"转换成功: {result if isinstance(result, str) else len(result)}个文件")
+        # 设置进度回调
+        def progress_callback(current, total):
+            print(f"\r进度: {current}/{total} ({current/total:.1%})", end='', flush=True)
+        converter.set_progress_callback(progress_callback)
         
-        # 示例3: 提取表格数据
-        print("\n提取PDF中的表格...")
-        csv_converter = PDFConverterFactory.get_converter('csv', 'example.pdf')
-        result = csv_converter.convert('output_tables', pages='all')
-        print(f"提取成功: {result if isinstance(result, str) else len(result)}个表格")
+        # 执行转换
+        print(f"开始转换: {args.input} -> {args.output} ({args.format})")
+        result = converter.convert(
+            args.output,
+            dpi=args.dpi,
+            start_page=args.start_page,
+            end_page=args.end_page,
+            preserve_formatting=args.preserve_formatting,
+            merge=args.merge,
+            quality=args.quality
+        )
         
-        # 查看所有支持的格式
-        print("\n支持的转换格式:")
-        for fmt, desc in PDFConverterFactory.get_supported_formats().items():
-            print(f"- {fmt}: {desc}")
+        print("\n转换成功!")
+        if isinstance(result, list):
+            print(f"生成文件 {len(result)} 个:")
+            for file in result:
+                print(f"  - {file}")
+        else:
+            print(f"输出文件: {result}")
             
     except Exception as e:
-        print(f"发生错误: {str(e)}")
+        print(f"\n转换失败: {e}", file=sys.stderr)
+        sys.exit(1)
+
+if __name__ == "__main__":
+    main()