Spaces:
Runtime error
Runtime error
File size: 4,558 Bytes
d28d608 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | """
PDF Processing Utilities
Handles PDF download and table extraction using pdfplumber
"""
import os
import requests
import pdfplumber
import pandas as pd
from typing import Optional, List
import logging
logger = logging.getLogger(__name__)
def download_pdf(url: str, save_dir: str = "downloads") -> str:
"""
Download PDF from URL to local directory
Args:
url: PDF file URL
save_dir: Directory to save the file
Returns:
Path to downloaded PDF file
"""
try:
# Ensure download directory exists
os.makedirs(save_dir, exist_ok=True)
# Extract filename from URL
filename = url.split("/")[-1]
if not filename.endswith(".pdf"):
filename = "document.pdf"
filepath = os.path.join(save_dir, filename)
# Download file
logger.info(f"Downloading PDF from {url}")
response = requests.get(url, timeout=30)
response.raise_for_status()
# Save to file
with open(filepath, "wb") as f:
f.write(response.content)
logger.info(f"PDF saved to {filepath}")
return filepath
except Exception as e:
logger.error(f"Error downloading PDF: {e}")
raise
def extract_tables(pdf_path: str, page_num: Optional[int] = None) -> List[pd.DataFrame]:
"""
Extract tables from PDF using pdfplumber
Args:
pdf_path: Path to PDF file
page_num: Specific page number (1-indexed), or None for all pages
Returns:
List of pandas DataFrames containing extracted tables
"""
try:
tables = []
with pdfplumber.open(pdf_path) as pdf:
# Determine which pages to process
if page_num is not None:
# Convert to 0-indexed
pages_to_process = [pdf.pages[page_num - 1]]
logger.info(f"Extracting tables from page {page_num}")
else:
pages_to_process = pdf.pages
logger.info(f"Extracting tables from all {len(pdf.pages)} pages")
# Extract tables from each page
for page_idx, page in enumerate(pages_to_process):
page_tables = page.extract_tables()
if page_tables:
for table_idx, table in enumerate(page_tables):
if table and len(table) > 0:
# Convert to DataFrame
# First row is typically headers
if len(table) > 1:
df = pd.DataFrame(table[1:], columns=table[0])
else:
df = pd.DataFrame(table)
# Clean column names
df.columns = [str(col).strip() if col else f"Column_{i}"
for i, col in enumerate(df.columns)]
# Remove empty rows
df = df.dropna(how='all')
tables.append(df)
logger.info(f"Extracted table {table_idx + 1} from page {page_idx + 1}: {df.shape}")
if not tables:
logger.warning("No tables found in PDF")
else:
logger.info(f"Total tables extracted: {len(tables)}")
return tables
except Exception as e:
logger.error(f"Error extracting tables from PDF: {e}")
raise
def extract_all_text(pdf_path: str, page_num: Optional[int] = None) -> str:
"""
Extract all text from PDF
Args:
pdf_path: Path to PDF file
page_num: Specific page number (1-indexed), or None for all pages
Returns:
Extracted text as string
"""
try:
text = []
with pdfplumber.open(pdf_path) as pdf:
if page_num is not None:
pages_to_process = [pdf.pages[page_num - 1]]
else:
pages_to_process = pdf.pages
for page in pages_to_process:
page_text = page.extract_text()
if page_text:
text.append(page_text)
return "\n".join(text)
except Exception as e:
logger.error(f"Error extracting text from PDF: {e}")
raise
|