wellfound-feedback / src /sheets_updater.py
Zoey7Web's picture
Upload 13 files
feef6eb verified
Raw
History Blame Contribute Delete
14.7 kB
"""
Google Sheets 回填模块
- 模糊匹配公司名
- 智能覆盖策略(优先级)
- 写入 AA/AB/AC/AD 列
所有配置值均通过 src.config 模块读取(支持运行时动态更新)。
"""
import logging
from dataclasses import dataclass, field
from datetime import datetime
from typing import Optional
import gspread
from google.oauth2.service_account import Credentials
from rapidfuzz import fuzz, process
from . import config as cfg
logger = logging.getLogger(__name__)
SCOPES = [
"https://www.googleapis.com/auth/spreadsheets",
"https://www.googleapis.com/auth/drive.file",
]
@dataclass
class MatchResult:
"""公司名匹配结果"""
status: str # exact | fuzzy | ambiguous | not_found
row_index: int = -1 # 1-based 行号(-1 表示未找到)
matched_name: str = "" # 表格中实际的公司名
score: float = 0.0 # 匹配分数 (0-100)
sheet_name: str = "" # 所在工作表名称(如 West、South 等)
candidates: list = field(default_factory=list) # 候选列表(ambiguous 时)
message: str = ""
@dataclass
class WriteResult:
"""写入结果"""
success: bool
row_index: int = -1
action: str = "" # written | skipped | overwritten
reason: str = ""
error: str = ""
def _get_priority(category: str) -> int:
"""获取分类优先级(数值越大优先级越高)"""
if category in cfg.HIGH_PRIORITY_TYPES:
return 2
if category in cfg.LOW_PRIORITY_TYPES:
return 1
return 0
class SheetsUpdater:
"""Google Sheets 更新器(跨多个区域工作表查询与写入)"""
def __init__(self):
self._client = None
self._spreadsheet = None
# 工作表缓存 {sheet_name: worksheet_object}
self._worksheets: dict = {}
# 公司列表:每条记录包含 (company_name, row_index, sheet_name)
self._company_list: list[tuple[str, int, str]] = []
self._initialized = False
def _get_credentials(self) -> Credentials:
"""获取 Google 凭证(每次调用读取最新配置)"""
sa_info = cfg.get_service_account_info()
if not sa_info:
raise ValueError(
"❌ 未配置 Google Service Account 凭证。\n"
"请在设置面板中粘贴 Service Account JSON 内容。"
)
return Credentials.from_service_account_info(sa_info, scopes=SCOPES)
def initialize(self) -> bool:
"""初始化连接并加载所有区域工作表的公司列表"""
try:
creds = self._get_credentials()
self._client = gspread.authorize(creds)
sheet_id = cfg.get_google_sheet_id()
if not sheet_id:
raise ValueError("❌ 未配置 Google Sheets ID。\n请在设置面板中填写 Google Sheet ID。")
self._spreadsheet = self._client.open_by_key(sheet_id)
self._worksheets = {}
# 加载所有区域工作表
sheet_names = cfg.get_google_sheet_names()
loaded = []
skipped = []
for name in sheet_names:
try:
ws = self._spreadsheet.worksheet(name)
self._worksheets[name] = ws
loaded.append(name)
except gspread.exceptions.WorksheetNotFound:
skipped.append(name)
logger.warning(f"⚠️ 工作表 '{name}' 不存在,已跳过")
if not self._worksheets:
raise ValueError(
f"❌ 未能加载任何区域工作表。\n"
f"已尝试: {sheet_names},请确认工作表名称正确。"
)
self._load_company_list()
self._initialized = True
logger.info(
f"✅ Google Sheets 连接成功,已加载工作表: {loaded},"
f"共 {len(self._company_list)} 家公司"
+ (f"(跳过不存在的工作表: {skipped})" if skipped else "")
)
return True
except Exception as e:
logger.error(f"Google Sheets 初始化失败: {e}")
raise
def _load_company_list(self):
"""加载所有区域工作表中的公司名(用于跨表模糊匹配)"""
indices = cfg.get_sheet_col_indices()
col_idx = indices["company"] + 1 # gspread 用 1-based
self._company_list = []
for sheet_name, ws in self._worksheets.items():
try:
col_values = ws.col_values(col_idx)
count = 0
for row_idx, name in enumerate(col_values, start=1):
name = name.strip()
if name and row_idx > 1: # 跳过表头
self._company_list.append((name, row_idx, sheet_name))
count += 1
logger.debug(f" [{sheet_name}] 已加载 {count} 家公司")
except Exception as e:
logger.warning(f"加载工作表 '{sheet_name}' 公司列表失败: {e}")
logger.debug(f"全部公司列表已加载: {len(self._company_list)} 条(跨 {len(self._worksheets)} 张工作表)")
def refresh_company_list(self):
"""刷新公司列表缓存"""
self._load_company_list()
def fuzzy_match_company(self, company_name: str) -> MatchResult:
"""模糊匹配公司名(跨所有区域工作表)"""
if not self._company_list:
return MatchResult(
status="not_found",
message="表格中无公司数据,请确认 Google Sheets 已正确配置"
)
# 提取名称列表,构建查找字典(同名公司保留首个——不同表可能重名)
names = [name for name, _, _ in self._company_list]
# 名称 -> (row, sheet_name),优先保留首次出现的
name_to_info: dict[str, tuple[int, str]] = {}
for name, row, sheet in self._company_list:
if name not in name_to_info:
name_to_info[name] = (row, sheet)
matches = process.extract(
company_name,
list(name_to_info.keys()),
scorer=fuzz.token_sort_ratio,
limit=5,
)
if not matches:
return MatchResult(status="not_found", message=f"未找到与 '{company_name}' 匹配的公司")
best_name, best_score, _ = matches[0]
best_row, best_sheet = name_to_info[best_name]
# 完全匹配
if company_name.lower() == best_name.lower():
return MatchResult(
status="exact",
row_index=best_row,
matched_name=best_name,
score=100.0,
sheet_name=best_sheet,
message=f"精确匹配: '{best_name}' ({best_sheet} 表, 第 {best_row} 行)"
)
threshold = cfg.get_fuzzy_threshold()
if best_score < threshold:
return MatchResult(
status="not_found",
message=f"未找到匹配项 (最高相似度 {best_score:.1f}% < 阈值 {threshold}%)"
)
# 收集所有达标候选
candidates = []
for name, score, _ in matches:
if score >= threshold:
row, sheet = name_to_info[name]
candidates.append({
"name": name,
"row": row,
"sheet": sheet,
"score": score
})
if len(candidates) >= 2:
gap = candidates[0]["score"] - candidates[1]["score"]
ambiguous_gap = cfg.get_fuzzy_ambiguous_gap()
if gap < ambiguous_gap:
return MatchResult(
status="ambiguous",
matched_name=best_name,
score=best_score,
sheet_name=best_sheet,
candidates=candidates[:5],
message=f"发现多个相似公司(差异 < {ambiguous_gap}%),需人工确认"
)
return MatchResult(
status="fuzzy",
row_index=best_row,
matched_name=best_name,
score=best_score,
sheet_name=best_sheet,
message=f"模糊匹配: '{best_name}' ({best_sheet} 表, 相似度 {best_score:.1f}%, 第 {best_row} 行)"
)
def _col_letter_to_1based(self, col_letter: str) -> int:
col = col_letter.upper().strip()
result = 0
for ch in col:
result = result * 26 + (ord(ch) - ord('A') + 1)
return result
def _get_worksheet(self, sheet_name: str):
"""根据工作表名获取 worksheet 对象,默认回退到第一个可用工作表"""
if sheet_name and sheet_name in self._worksheets:
return self._worksheets[sheet_name]
# 回退:使用第一个可用工作表
if self._worksheets:
fallback = next(iter(self._worksheets))
logger.warning(f"工作表 '{sheet_name}' 未找到,回退到 '{fallback}'")
return self._worksheets[fallback]
raise RuntimeError("没有可用的工作表,请重新初始化连接")
def _get_cell_value(self, row: int, col_letter: str, sheet_name: str = "") -> str:
try:
ws = self._get_worksheet(sheet_name)
col_idx = self._col_letter_to_1based(col_letter)
cell = ws.cell(row, col_idx)
return (cell.value or "").strip()
except Exception as e:
logger.warning(f"读取单元格失败 (sheet={sheet_name}, row={row}, col={col_letter}): {e}")
return ""
def _write_cells(self, row: int, data: dict, sheet_name: str = "") -> None:
"""批量写入单元格 {col_letter: value},写入指定工作表"""
ws = self._get_worksheet(sheet_name)
updates = []
for col_letter, value in data.items():
col_idx = self._col_letter_to_1based(col_letter)
cell = gspread.utils.rowcol_to_a1(row, col_idx)
updates.append({
"range": cell,
"values": [[str(value) if value is not None else ""]]
})
if updates:
ws.batch_update(updates, value_input_option="USER_ENTERED")
def write_email_data(
self,
row_index: int,
category: str,
body_text: str,
date_str: str,
force_overwrite: bool = False,
sheet_name: str = ""
) -> WriteResult:
"""按优先级策略写入邮件数据到指定工作表"""
try:
cfg_dict = cfg.get_config_dict()
existing_type = self._get_cell_value(
row_index, cfg_dict["sheet_email_type_col"], sheet_name
)
if existing_type and not force_overwrite:
new_priority = _get_priority(category)
existing_priority = _get_priority(existing_type)
if new_priority > existing_priority:
action_reason = f"覆盖({existing_type} -> {category},新邮件优先级更高)"
elif new_priority <= existing_priority and existing_type == category:
return WriteResult(
success=True, row_index=row_index,
action="skipped",
reason=f"相同类型 '{category}' 已存在,保留原数据"
)
elif new_priority < existing_priority:
return WriteResult(
success=True, row_index=row_index,
action="skipped",
reason=f"已有高优先级数据 '{existing_type}',新邮件 '{category}' 被忽略"
)
else:
return WriteResult(
success=True, row_index=row_index,
action="skipped",
reason=f"同优先级数据已存在,保留原数据 '{existing_type}'"
)
else:
action_reason = "首次写入" if not existing_type else "强制覆盖(人工确认)"
body_truncated = body_text[:cfg.get_email_body_max_chars()] if body_text else ""
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
write_data = {
cfg_dict["sheet_email_type_col"]: category,
cfg_dict["sheet_email_body_col"]: body_truncated,
cfg_dict["sheet_email_date_col"]: date_str,
cfg_dict["sheet_timestamp_col"]: timestamp,
}
self._write_cells(row_index, write_data, sheet_name)
target = f"[{sheet_name}] " if sheet_name else ""
logger.info(f"✅ {target}Row {row_index}: 写入成功 [{category}] - {action_reason}")
return WriteResult(
success=True,
row_index=row_index,
action="overwritten" if existing_type else "written",
reason=action_reason
)
except Exception as e:
logger.error(f"写入 Google Sheets 失败 (sheet={sheet_name}, row={row_index}): {e}", exc_info=True)
return WriteResult(
success=False,
row_index=row_index,
action="error",
error=str(e)
)
def is_connected(self) -> bool:
return self._initialized and bool(self._worksheets)
def get_company_count(self) -> int:
return len(self._company_list)
def get_candidates_for_display(self, candidates: list) -> list[list]:
return [
[c["name"], c.get("sheet", ""), f"{c['score']:.1f}%", f"第 {c['row']} 行"]
for c in candidates
]
_updater_instance = None
def get_sheets_updater(auto_init: bool = True):
"""获取 SheetsUpdater 单例(惰性初始化)"""
global _updater_instance
if _updater_instance is None:
_updater_instance = SheetsUpdater()
if auto_init and not _updater_instance.is_connected():
try:
_updater_instance.initialize()
except Exception as e:
logger.warning(f"Sheets 自动初始化失败(将在使用时重试): {e}")
return _updater_instance
def reset_sheets_updater():
"""重置单例(配置变更后调用,下次使用时会重新初始化)"""
global _updater_instance
_updater_instance = None