Spaces:
Runtime error
Runtime error
| """ | |
| Excel Handler - Read/write Excel files with checkpoint/resume support. | |
| Handles: | |
| - Loading Excel files with column detection | |
| - Incremental save with checkpoint tracking | |
| - Resume from last processed row | |
| - Column mapping (Wellfound export format) | |
| """ | |
| import json | |
| import os | |
| import threading | |
| from pathlib import Path | |
| from typing import Optional, Dict, Any, List, Tuple | |
| import pandas as pd | |
| import openpyxl | |
| from openpyxl.utils import get_column_letter | |
| class ExcelHandler: | |
| """Handles Excel file operations with checkpoint support.""" | |
| # Wellfound export column mapping | |
| COLUMN_MAP = { | |
| "company_name": "Company Name", | |
| "location": "Location", | |
| "internal_link": "Internal Link", | |
| "external_link": "External Link", | |
| "valuation": "Valuation", | |
| "rounds": "Rounds", | |
| "series": "Series", | |
| "total_raised": "Total Raised", | |
| "location_apply": "location.apply", | |
| "state_apply": "state.apply", | |
| "region": "Region", | |
| "contact": "contact", | |
| } | |
| def __init__(self, checkpoint_dir: str = "data/checkpoints"): | |
| self.checkpoint_dir = Path(checkpoint_dir) | |
| self.checkpoint_dir.mkdir(parents=True, exist_ok=True) | |
| self._lock = threading.Lock() | |
| def load(self, file_path: str) -> pd.DataFrame: | |
| """Load Excel file into DataFrame.""" | |
| df = pd.read_excel(file_path, engine="openpyxl") | |
| # Ensure all expected columns exist | |
| for col in self.COLUMN_MAP.values(): | |
| if col not in df.columns: | |
| df[col] = None | |
| return df | |
| def get_checkpoint(self, file_hash: str) -> Optional[Dict[str, Any]]: | |
| """Get checkpoint state for a file.""" | |
| checkpoint_file = self.checkpoint_dir / f"{file_hash}.json" | |
| if checkpoint_file.exists(): | |
| with open(checkpoint_file, "r", encoding="utf-8") as f: | |
| return json.load(f) | |
| return None | |
| def save_checkpoint(self, file_hash: str, state: Dict[str, Any]): | |
| """Save checkpoint state.""" | |
| checkpoint_file = self.checkpoint_dir / f"{file_hash}.json" | |
| with self._lock: | |
| with open(checkpoint_file, "w", encoding="utf-8") as f: | |
| json.dump(state, f, ensure_ascii=False, indent=2) | |
| def save_row(self, file_path: str, row_idx: int, updates: Dict[str, Any]): | |
| """Save a single row's updates to Excel (incremental write).""" | |
| with self._lock: | |
| wb = openpyxl.load_workbook(file_path) | |
| ws = wb.active | |
| # Find column indices | |
| col_indices = {} | |
| for cell in ws[1]: | |
| if cell.value: | |
| col_indices[str(cell.value)] = cell.column | |
| excel_row = row_idx + 2 # +2 for header and 0-index | |
| for col_name, value in updates.items(): | |
| if col_name in col_indices: | |
| col_idx = col_indices[col_name] | |
| ws.cell(row=excel_row, column=col_idx, value=value) | |
| wb.save(file_path) | |
| wb.close() | |
| def save_all_rows(self, file_path: str, df: pd.DataFrame): | |
| """Save entire DataFrame to Excel.""" | |
| with self._lock: | |
| df.to_excel(file_path, index=False, engine="openpyxl") | |
| def get_resume_index(self, file_hash: str) -> int: | |
| """Get the row index to resume from.""" | |
| checkpoint = self.get_checkpoint(file_hash) | |
| if checkpoint and "last_processed_index" in checkpoint: | |
| return checkpoint["last_processed_index"] + 1 | |
| return 0 | |
| def create_output_copy(self, source_path: str, output_dir: str = "data/results") -> str: | |
| """Create a working copy of the Excel file for incremental saves.""" | |
| output_dir = Path(output_dir) | |
| output_dir.mkdir(parents=True, exist_ok=True) | |
| source_name = Path(source_path).stem | |
| output_path = output_dir / f"{source_name}_processed.xlsx" | |
| df = self.load(source_path) | |
| df.to_excel(str(output_path), index=False, engine="openpyxl") | |
| return str(output_path) | |
| def compute_file_hash(file_path: str) -> str: | |
| """Compute a simple hash for the file.""" | |
| import hashlib | |
| stat = os.stat(file_path) | |
| name = Path(file_path).name | |
| raw = f"{name}_{stat.st_size}_{stat.st_mtime}" | |
| return hashlib.md5(raw.encode()).hexdigest()[:12] | |