| """ |
| 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.""" |
|
|
| |
| 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") |
| |
| 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 |
|
|
| |
| col_indices = {} |
| for cell in ws[1]: |
| if cell.value: |
| col_indices[str(cell.value)] = cell.column |
|
|
| excel_row = row_idx + 2 |
|
|
| 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) |
|
|
| @staticmethod |
| 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] |
|
|
