Update app.py

app.py

@@ -31,37 +31,13 @@ TEMP_DIR.mkdir(exist_ok=True)
 user_sessions = {}
 
 class PDFProcessor:
-    """Handle PDF splitting with qpdf/pikepdf"""
-    
-    @staticmethod
-    def estimate_pages_for_size(pdf, total_pages: int, target_size_bytes: int) -> int:
-        """
-        Estimate how many pages fit in the target size
-        """
-        # Get approximate file size
-        temp_file = Path("temp_estimate.pdf")
-        try:
-            # Save the entire PDF temporarily to get its size
-            pdf.save(temp_file)
-            total_size = temp_file.stat().st_size
-            temp_file.unlink()
-            
-            # Calculate average page size
-            avg_page_size = total_size / total_pages if total_pages > 0 else total_size
-            
-            # Estimate pages that fit in target size (with 10% safety margin)
-            estimated_pages = int((target_size_bytes * 0.9) / avg_page_size)
-            
-            return max(1, estimated_pages)  # At least 1 page
-            
-        except Exception as e:
-            logger.error(f"Error estimating page size: {e}")
-            return max(1, int(total_pages / 10))  # Fallback to 10% of pages
+    """Handle PDF splitting with qpdf/pikepdf - with corruption handling"""
     
     @staticmethod
     def split_pdf_by_size(input_path: Path, output_dir: Path, progress_callback=None) -> Tuple[List[Path], dict]:
         """
         Split PDF into segments of approximately 4.5MB, discarding any over 5MB
+        Handles corrupted PDFs by attempting recovery
         """
         kept_files = []
         discarded_count = 0
@@ -79,18 +55,71 @@ class PDFProcessor:
             # Get original file size
             stats["original_size_mb"] = input_path.stat().st_size / 1024 / 1024
             
-            # Open PDF with pikepdf
-            with pikepdf.open(input_path) as pdf:
+            # First attempt: Try to open with recovery and stream decoding disabled
+            pdf = None
+            try:
+                # Open PDF with recovery mode and suppress stream errors
+                pdf = pikepdf.open(
+                    input_path,
+                    suppress_warnings=True,
+                    attempt_recovery=True
+                )
+            except pikepdf._qpdf.DataDecodingError as e:
+                logger.warning(f"Initial open failed, attempting repair: {e}")
+                
+                # Second attempt: Create a repaired copy first
+                repaired_path = output_dir.parent / "repaired_temp.pdf"
+                try:
+                    # Use pikepdf to create a repaired version
+                    with pikepdf.open(input_path, suppress_warnings=True, attempt_recovery=True) as damaged_pdf:
+                        # Save with recompression to fix stream errors
+                        damaged_pdf.save(
+                            repaired_path,
+                            compress_streams=False,  # Disable compression first
+                            decode_level=pikepdf.StreamDecodeLevel.all,  # Decode all streams
+                            object_stream_mode=pikepdf.ObjectStreamMode.disable,  # Disable object streams
+                            normalize_content=True,  # Normalize content streams
+                            linearize=False
+                        )
+                    
+                    # Now open the repaired version
+                    pdf = pikepdf.open(repaired_path, suppress_warnings=True)
+                    
+                    # Clean up repaired file after opening
+                    input_path = repaired_path
+                    
+                except Exception as repair_error:
+                    logger.error(f"Repair attempt failed: {repair_error}")
+                    
+                    # Third attempt: Try with qpdf command line if available
+                    import subprocess
+                    try:
+                        repaired_path = output_dir.parent / "qpdf_repaired.pdf"
+                        subprocess.run(
+                            ["qpdf", "--replace-input", "--stream-data=uncompress", 
+                             str(input_path), str(repaired_path)],
+                            check=True,
+                            capture_output=True
+                        )
+                        pdf = pikepdf.open(repaired_path, suppress_warnings=True)
+                        input_path = repaired_path
+                    except (subprocess.CalledProcessError, FileNotFoundError):
+                        raise Exception("PDF is severely corrupted and cannot be repaired")
+            
+            if pdf is None:
+                raise Exception("Failed to open PDF after all recovery attempts")
+            
+            with pdf:
                 total_pages = len(pdf.pages)
                 stats["total_pages"] = total_pages
                 
                 if total_pages == 0:
                     return kept_files, stats
                 
-                # Initial estimate of pages per segment
-                pages_per_segment = PDFProcessor.estimate_pages_for_size(
-                    pdf, total_pages, TARGET_SEGMENT_SIZE_BYTES
-                )
+                # Estimate pages per segment
+                file_size = input_path.stat().st_size
+                avg_page_size = file_size / total_pages if total_pages > 0 else file_size
+                pages_per_segment = max(1, int(TARGET_SEGMENT_SIZE_BYTES * 0.8 / avg_page_size))
                 
                 segment_num = 0
                 page_start = 0
@@ -98,99 +127,113 @@ class PDFProcessor:
                 max_retries = 3
                 
                 while page_start < total_pages:
-                    # Calculate page range for this segment
                     page_end = min(page_start + pages_per_segment, total_pages)
                     
-                    # Update progress
                     if progress_callback:
                         progress = (page_start / total_pages)
                         progress_callback(progress, f"Processing pages {page_start+1}-{page_end} of {total_pages}...")
                     
-                    # Create segment
                     segment_num += 1
                     segment_filename = f"segment_{segment_num:03d}_p{page_start+1}-{page_end}.pdf"
                     segment_path = output_dir / segment_filename
                     
-                    # Create new PDF with selected pages
-                    segment_pdf = pikepdf.new()
-                    for page_num in range(page_start, page_end):
-                        segment_pdf.pages.append(pdf.pages[page_num])
-                    
-                    # Save with compression to minimize size
-                    segment_pdf.save(
-                        segment_path,
-                        compress_streams=True,
-                        stream_decode_level=pikepdf.StreamDecodeLevel.none,
-                        object_stream_mode=pikepdf.ObjectStreamMode.generate,
-                        linearize=True,
-                        recompress_flate=True
-                    )
-                    
-                    # Check segment size
-                    segment_size = segment_path.stat().st_size
-                    segment_size_mb = segment_size / 1024 / 1024
-                    
-                    logger.info(f"Segment {segment_num}: {segment_size_mb:.2f} MB ({page_end - page_start} pages)")
-                    
-                    if segment_size <= MAX_ALLOWED_SIZE_BYTES:
-                        # File is under 5MB limit - keep it
-                        kept_files.append(segment_path)
-                        stats["segments_created"] += 1
-                        stats["total_output_size_mb"] += segment_size_mb
+                    try:
+                        # Create new PDF with selected pages
+                        segment_pdf = pikepdf.new()
                         
-                        # Track largest and smallest segments
-                        stats["largest_segment_mb"] = max(stats["largest_segment_mb"], segment_size_mb)
-                        stats["smallest_segment_mb"] = min(stats["smallest_segment_mb"], segment_size_mb)
+                        # Copy pages with error handling
+                        for page_num in range(page_start, page_end):
+                            try:
+                                # Clone the page to avoid stream errors
+                                page = pdf.pages[page_num]
+                                segment_pdf.pages.append(page)
+                            except Exception as page_error:
+                                logger.warning(f"Error copying page {page_num+1}: {page_error}")
+                                # Skip corrupted pages
+                                continue
                         
-                        # Move to next segment
-                        page_start = page_end
-                        retry_count = 0  # Reset retry count for next segment
+                        # If no pages were successfully added, skip this segment
+                        if len(segment_pdf.pages) == 0:
+                            logger.warning(f"Segment {segment_num} has no valid pages, skipping")
+                            page_start = page_end
+                            continue
                         
-                        # Adjust pages per segment based on actual size
-                        if segment_size_mb < 4.0 and pages_per_segment < total_pages:
-                            # Segment is too small, try more pages next time
-                            pages_per_segment = min(pages_per_segment + 1, total_pages - page_end)
-                        elif segment_size_mb > 4.8:
-                            # Segment is getting close to limit, use fewer pages
-                            pages_per_segment = max(1, pages_per_segment - 1)
-                            
-                    else:
-                        # File exceeds 5MB limit
-                        logger.warning(f"Segment {segment_num} too large ({segment_size_mb:.2f} MB)")
+                        # Save with safe compression settings
+                        segment_pdf.save(
+                            segment_path,
+                            compress_streams=True,
+                            stream_decode_level=pikepdf.StreamDecodeLevel.specialized,  # Use specialized decoding
+                            object_stream_mode=pikepdf.ObjectStreamMode.generate,
+                            normalize_content=True,  # Normalize to fix issues
+                            linearize=False,  # Don't linearize to avoid issues
+                            recompress_flate=False  # Don't recompress to avoid corruption
+                        )
                         
-                        if page_end - page_start == 1:
-                            # Single page is over 5MB - discard and move on
-                            logger.warning(f"Single page {page_start+1} exceeds 5MB limit - discarding")
-                            segment_path.unlink()  # Delete the file
-                            stats["segments_discarded"] += 1
-                            discarded_count += 1
-                            page_start = page_end  # Move to next page
+                    except Exception as save_error:
+                        logger.error(f"Error saving segment {segment_num}: {save_error}")
+                        # Try saving without compression
+                        try:
+                            segment_pdf.save(
+                                segment_path,
+                                compress_streams=False,
+                                object_stream_mode=pikepdf.ObjectStreamMode.disable
+                            )
+                        except:
+                            logger.error(f"Failed to save segment {segment_num} even without compression")
+                            page_start = page_end
+                            continue
+                    
+                    # Check segment size
+                    if segment_path.exists():
+                        segment_size = segment_path.stat().st_size
+                        segment_size_mb = segment_size / 1024 / 1024
+                        
+                        logger.info(f"Segment {segment_num}: {segment_size_mb:.2f} MB")
+                        
+                        if segment_size <= MAX_ALLOWED_SIZE_BYTES:
+                            kept_files.append(segment_path)
+                            stats["segments_created"] += 1
+                            stats["total_output_size_mb"] += segment_size_mb
+                            stats["largest_segment_mb"] = max(stats["largest_segment_mb"], segment_size_mb)
+                            stats["smallest_segment_mb"] = min(stats["smallest_segment_mb"], segment_size_mb)
+                            page_start = page_end
                             retry_count = 0
                         else:
-                            # Multiple pages - try with fewer pages
-                            segment_path.unlink()  # Delete the oversized file
+                            # File exceeds 5MB limit
+                            logger.warning(f"Segment {segment_num} too large ({segment_size_mb:.2f} MB)")
                             
-                            if retry_count < max_retries:
-                                # Reduce pages by half and retry
-                                pages_per_segment = max(1, (page_end - page_start) // 2)
-                                retry_count += 1
-                                segment_num -= 1  # Reuse segment number
-                                logger.info(f"Retrying with {pages_per_segment} pages")
+                            if page_end - page_start == 1:
+                                # Single page is over 5MB - discard
+                                segment_path.unlink()
+                                stats["segments_discarded"] += 1
+                                page_start = page_end
                             else:
-                                # Too many retries, try single pages
-                                pages_per_segment = 1
-                                retry_count = 0
+                                # Try with fewer pages
+                                segment_path.unlink()
+                                pages_per_segment = max(1, (page_end - page_start) // 2)
                                 segment_num -= 1
                 
-                # Clean up stats
                 if stats["smallest_segment_mb"] == float('inf'):
                     stats["smallest_segment_mb"] = 0
                 
                 if progress_callback:
                     progress_callback(1.0, "Splitting complete!")
+            
+            # Clean up temporary repaired files if they exist
+            for temp_file in output_dir.parent.glob("*repaired*.pdf"):
+                try:
+                    temp_file.unlink()
+                except:
+                    pass
                     
         except Exception as e:
             logger.error(f"Error splitting PDF: {str(e)}")
+            # Clean up any temporary files
+            for temp_file in output_dir.parent.glob("*repaired*.pdf"):
+                try:
+                    temp_file.unlink()
+                except:
+                    pass
             raise
         
         return kept_files, stats