Update app.py

app.py

@@ -31,16 +31,15 @@ TEMP_DIR.mkdir(exist_ok=True)
 user_sessions = {}
 
 class PDFProcessor:
-    """Handle PDF splitting with qpdf/pikepdf - with corruption handling"""
+    """Handle PDF splitting with qpdf/pikepdf - using incremental size checking like bash script"""
     
     @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
+        Uses the same incremental approach as the bash script
         """
         kept_files = []
-        discarded_count = 0
         stats = {
             "total_pages": 0,
             "segments_created": 0,
@@ -55,181 +54,146 @@ class PDFProcessor:
             # Get original file size
             stats["original_size_mb"] = input_path.stat().st_size / 1024 / 1024
             
-            # 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:
+            # Open PDF with pikepdf
+            with pikepdf.open(input_path, suppress_warnings=True, attempt_recovery=True) as pdf:
                 total_pages = len(pdf.pages)
                 stats["total_pages"] = total_pages
                 
                 if total_pages == 0:
                     return kept_files, stats
                 
-                # 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
-                retry_count = 0
-                max_retries = 3
+                start_page = 0
+                part = 1
                 
-                while page_start < total_pages:
-                    page_end = min(page_start + pages_per_segment, total_pages)
+                while start_page < total_pages:
+                    # Start with a single page
+                    end_page = start_page
+                    temp_segment = None
+                    last_good_segment = None
+                    last_good_end = start_page
                     
+                    # Update progress
                     if progress_callback:
-                        progress = (page_start / total_pages)
-                        progress_callback(progress, f"Processing pages {page_start+1}-{page_end} of {total_pages}...")
+                        progress = (start_page / total_pages)
+                        progress_callback(progress, f"Processing segment {part}, starting at page {start_page + 1}...")
                     
-                    segment_num += 1
-                    segment_filename = f"segment_{segment_num:03d}_p{page_start+1}-{page_end}.pdf"
-                    segment_path = output_dir / segment_filename
-                    
-                    try:
-                        # Create new PDF with selected pages
-                        segment_pdf = pikepdf.new()
-                        
-                        # 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
-                        
-                        # 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
-                        
-                        # 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
-                        )
+                    # Keep adding pages until we exceed the size limit
+                    while end_page < total_pages:
+                        # Create temporary segment with pages from start_page to end_page (inclusive)
+                        temp_filename = f"temp_segment_{part}.pdf"
+                        temp_path = output_dir / temp_filename
                         
-                    except Exception as save_error:
-                        logger.error(f"Error saving segment {segment_num}: {save_error}")
-                        # Try saving without compression
                         try:
+                            # Create new PDF with selected pages
+                            segment_pdf = pikepdf.new()
+                            
+                            # Add pages from start_page to end_page (inclusive)
+                            for page_num in range(start_page, end_page + 1):
+                                segment_pdf.pages.append(pdf.pages[page_num])
+                            
+                            # Save with compression
                             segment_pdf.save(
-                                segment_path,
-                                compress_streams=False,
-                                object_stream_mode=pikepdf.ObjectStreamMode.disable
+                                temp_path,
+                                compress_streams=True,
+                                object_stream_mode=pikepdf.ObjectStreamMode.generate,
+                                linearize=False  # Don't linearize to save time during testing
                             )
-                        except:
-                            logger.error(f"Failed to save segment {segment_num} even without compression")
-                            page_start = page_end
-                            continue
+                            
+                            # Check file size
+                            segment_size = temp_path.stat().st_size
+                            segment_size_mb = segment_size / 1024 / 1024
+                            
+                            logger.debug(f"Testing segment {part}: pages {start_page+1}-{end_page+1}, size: {segment_size_mb:.2f} MB")
+                            
+                            if segment_size < TARGET_SEGMENT_SIZE_BYTES:
+                                # Still under target size, keep this as last good and try adding more pages
+                                if last_good_segment and last_good_segment.exists():
+                                    last_good_segment.unlink()  # Delete previous good segment
+                                last_good_segment = temp_path
+                                last_good_end = end_page
+                                
+                                # If we're at the last page, this is our final segment
+                                if end_page == total_pages - 1:
+                                    break
+                                
+                                # Try adding one more page
+                                end_page += 1
+                                
+                            elif segment_size <= MAX_ALLOWED_SIZE_BYTES:
+                                # Between 4.5MB and 5MB - this is acceptable, use it
+                                if last_good_segment and last_good_segment.exists():
+                                    last_good_segment.unlink()
+                                last_good_segment = temp_path
+                                last_good_end = end_page
+                                break  # Stop here, we found a good size
+                                
+                            else:
+                                # Over 5MB limit
+                                temp_path.unlink()  # Delete oversized segment
+                                
+                                if end_page == start_page:
+                                    # Single page is over 5MB - discard it
+                                    logger.warning(f"Single page {start_page+1} exceeds 5MB limit - discarding")
+                                    stats["segments_discarded"] += 1
+                                    last_good_end = start_page  # Move past this page
+                                    break
+                                else:
+                                    # Multiple pages - use the last good segment
+                                    break
+                            
+                        except Exception as e:
+                            logger.error(f"Error creating segment: {e}")
+                            if temp_path and temp_path.exists():
+                                temp_path.unlink()
+                            break
                     
-                    # Check segment size
-                    if segment_path.exists():
-                        segment_size = segment_path.stat().st_size
-                        segment_size_mb = segment_size / 1024 / 1024
+                    # Save the final segment for this part
+                    if last_good_segment and last_good_segment.exists():
+                        # Rename to final name
+                        final_filename = f"segment_{part:03d}_p{start_page+1}-{last_good_end+1}.pdf"
+                        final_path = output_dir / final_filename
+                        last_good_segment.rename(final_path)
                         
-                        logger.info(f"Segment {segment_num}: {segment_size_mb:.2f} MB")
+                        # Check final size and add to kept files
+                        final_size = final_path.stat().st_size
+                        final_size_mb = final_size / 1024 / 1024
                         
-                        if segment_size <= MAX_ALLOWED_SIZE_BYTES:
-                            kept_files.append(segment_path)
+                        if final_size <= MAX_ALLOWED_SIZE_BYTES:
+                            kept_files.append(final_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:
-                            # File exceeds 5MB limit
-                            logger.warning(f"Segment {segment_num} too large ({segment_size_mb:.2f} MB)")
+                            stats["total_output_size_mb"] += final_size_mb
+                            stats["largest_segment_mb"] = max(stats["largest_segment_mb"], final_size_mb)
+                            stats["smallest_segment_mb"] = min(stats["smallest_segment_mb"], final_size_mb)
                             
-                            if page_end - page_start == 1:
-                                # Single page is over 5MB - discard
-                                segment_path.unlink()
-                                stats["segments_discarded"] += 1
-                                page_start = page_end
-                            else:
-                                # Try with fewer pages
-                                segment_path.unlink()
-                                pages_per_segment = max(1, (page_end - page_start) // 2)
-                                segment_num -= 1
+                            logger.info(f"Created segment {part}: {final_size_mb:.2f} MB (pages {start_page+1}-{last_good_end+1})")
+                        else:
+                            # Should not happen, but just in case
+                            final_path.unlink()
+                            stats["segments_discarded"] += 1
+                            logger.warning(f"Final segment {part} exceeded 5MB limit after rename")
+                    
+                    # Move to next segment
+                    start_page = last_good_end + 1
+                    part += 1
+                    
+                    # Clean up any remaining temp files
+                    for temp_file in output_dir.glob("temp_segment_*.pdf"):
+                        try:
+                            temp_file.unlink()
+                        except:
+                            pass
                 
+                # Final cleanup
                 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"):
+            # Clean up temp files on error
+            for temp_file in output_dir.glob("temp_segment_*.pdf"):
                 try:
                     temp_file.unlink()
                 except: