Upload edgar_client.py

edgar_client.py

@@ -12,6 +12,8 @@ import time
 import threading
 from functools import lru_cache
 from datetime import datetime, timedelta
+import re
+import difflib
 
 
 class EdgarDataClient:
@@ -26,10 +28,40 @@ class EdgarDataClient:
     _rate_limit_lock = threading.Lock()
     _min_request_interval = 0.11  # 110ms between requests (9 req/sec, safe margin)
     
+    # 新增：公司索引（加速搜索，避免每次遍历全量数据）
+    _by_ticker = None  # ticker -> company info
+    _by_title = None  # title (lowercase) -> company info
+    _by_title_norm = None  # normalized title -> company info
+    _all_keys = None  # 用于模糊匹配的所有key列表
+    _index_built_time = None
+    _index_ttl = 3600  # 1 hour
+    
+    # 新增：常见别名映射（提升搜索智能性）
+    _alias_map = {
+        "google": "alphabet inc",
+        "alphabet": "alphabet inc",
+        "facebook": "meta platforms, inc.",
+        "meta": "meta platforms, inc.",
+        "amazon": "amazon.com, inc.",
+        "apple": "apple inc.",
+        "microsoft": "microsoft corporation",
+        "netflix": "netflix, inc.",
+        "nvidia": "nvidia corporation",
+        "tesla": "tesla, inc.",
+        "adobe": "adobe inc.",
+        "oracle": "oracle corporation",
+        "ibm": "international business machines corporation",
+        "paypal": "paypal holdings, inc.",
+        "shopify": "shopify inc.",
+    }
+    
     def __init__(self, user_agent="Juntao Peng Financial Report Metrics App (jtyxabc@gmail.com)"):
         """Initialize EDGAR client with connection pooling and timeout"""
         self.user_agent = user_agent
         
+        # 新增：实例级搜索缓存（进一步减少重复搜索开销）
+        self._search_cache = {}
+        
         # Configure requests session with connection pooling
         self.session = requests.Session()
         
@@ -131,76 +163,178 @@ class EdgarDataClient:
                 time.sleep(sleep_time)
             
             EdgarDataClient._last_request_time = time.time()
-        
-    def search_company_by_name(self, company_name):
-        """Search company CIK by company name with caching and optimized ticker matching"""
-        try:
-            # Check cache first
-            with self._cache_lock:
-                current_time = time.time()
+    
+    def _normalize_text(self, s: str) -> str:
+        """规范化文本：用于提升匹配准确度"""
+        if not s:
+            return ""
+        s = s.lower().strip()
+        s = s.replace("&", " and ")
+        s = re.sub(r"[.,()\-_/]", " ", s)
+        s = re.sub(r"\s+", " ", s)
+        # 移除常见后缀词
+        stopwords = {"inc", "inc.", "incorporated", "corp", "corporation", "co", "company", "plc", "ltd", "llc", "the"}
+        tokens = [t for t in s.split() if t not in stopwords]
+        return " ".join(tokens).strip()
+    
+    def _ensure_company_index(self):
+        """确保公司索引已构建（按需构建或过期重建）"""
+        with self._cache_lock:
+            current_time = time.time()
+            
+            # 若 company_tickers 缓存不存在或已过期，先刷新
+            if (EdgarDataClient._company_tickers_cache is None or
+                EdgarDataClient._company_tickers_cache_time is None or
+                current_time - EdgarDataClient._company_tickers_cache_time >= self._company_tickers_cache_ttl):
+                # 拉取并更新 company_tickers 缓存
+                self._rate_limit()
+                url = "https://www.sec.gov/files/company_tickers.json"
+                headers = {"User-Agent": self.user_agent}
+                response = self.session.get(url, headers=headers, timeout=self.timeout)
+                response.raise_for_status()
+                companies = response.json()
+                EdgarDataClient._company_tickers_cache = companies
+                EdgarDataClient._company_tickers_cache_time = current_time
+            else:
+                companies = EdgarDataClient._company_tickers_cache
+            
+            # 若索引不存在或已过期，则重建索引
+            if (EdgarDataClient._by_ticker is None or
+                EdgarDataClient._by_title is None or
+                EdgarDataClient._by_title_norm is None or
+                EdgarDataClient._all_keys is None or
+                EdgarDataClient._index_built_time is None or
+                current_time - EdgarDataClient._index_built_time >= EdgarDataClient._index_ttl):
                 
-                # If cache is valid, use it
-                if (EdgarDataClient._company_tickers_cache is not None and 
-                    EdgarDataClient._company_tickers_cache_time is not None and
-                    current_time - EdgarDataClient._company_tickers_cache_time < self._company_tickers_cache_ttl):
-                    companies = EdgarDataClient._company_tickers_cache
-                else:
-                    # Cache miss or expired, fetch new data
-                    self._rate_limit()
-                    url = "https://www.sec.gov/files/company_tickers.json"
-                    headers = {"User-Agent": self.user_agent}
-                    
-                    response = self.session.get(url, headers=headers, timeout=self.timeout)
-                    response.raise_for_status()
+                by_ticker = {}
+                by_title = {}
+                by_title_norm = {}
+                all_keys = []
+                
+                for _, company in companies.items():
+                    title = company.get("title", "")
+                    ticker = company.get("ticker", "")
+                    cik_str = str(company.get("cik_str", "")).zfill(10)
                     
-                    companies = response.json()
+                    title_lower = title.lower()
+                    ticker_lower = ticker.lower()
+                    title_norm = self._normalize_text(title)
                     
-                    # Update cache
-                    EdgarDataClient._company_tickers_cache = companies
-                    EdgarDataClient._company_tickers_cache_time = current_time
-            
-            # ✅ OPTIMIZATION 1: Prioritize exact ticker match (fastest path)
-            search_name_upper = company_name.upper().strip()
-            for _, company in companies.items():
-                if company["ticker"].upper() == search_name_upper:
-                    # Exact ticker match - return immediately
-                    return {
-                        "cik": str(company["cik_str"]).zfill(10),
-                        "name": company["title"],
-                        "ticker": company["ticker"]
-                    }
-            
-            # ✅ OPTIMIZATION 2: Search for matching company names
-            matches = []
-            exact_matches = []
-            search_name_lower = company_name.lower()
-            
-            for _, company in companies.items():
-                company_title = company["title"].lower()
-                ticker_lower = company["ticker"].lower()
+                    # 构建索引：ticker、title、normalized title
+                    if ticker_lower:
+                        by_ticker[ticker_lower] = {"cik": cik_str, "name": title, "ticker": ticker}
+                        all_keys.append(ticker_lower)
+                    if title_lower:
+                        by_title[title_lower] = {"cik": cik_str, "name": title, "ticker": ticker}
+                    if title_norm:
+                        by_title_norm[title_norm] = {"cik": cik_str, "name": title, "ticker": ticker}
+                        all_keys.append(title_norm)
                 
-                # Exact match
-                if search_name_lower == company_title:
-                    exact_matches.append({
-                        "cik": str(company["cik_str"]).zfill(10),
-                        "name": company["title"],
-                        "ticker": company["ticker"]
-                    })
-                # Partial match (name or ticker contains search term)
-                elif search_name_lower in company_title or search_name_lower in ticker_lower:
-                    matches.append({
-                        "cik": str(company["cik_str"]).zfill(10),
-                        "name": company["title"],
-                        "ticker": company["ticker"]
-                    })
-            
-            # Return exact match first, then partial match
-            if exact_matches:
-                return exact_matches[0]
-            elif matches:
-                return matches[0]
-            else:
-                return None
+                EdgarDataClient._by_ticker = by_ticker
+                EdgarDataClient._by_title = by_title
+                EdgarDataClient._by_title_norm = by_title_norm
+                EdgarDataClient._all_keys = all_keys
+                EdgarDataClient._index_built_time = current_time
+        
+    def search_company_by_name(self, company_name):
+        """Search company CIK by company name with caching and optimized ticker matching"""
+        try:
+            # 实例级缓存命中检查（按规范化后的query）
+            norm_query = self._normalize_text(company_name)
+            cache_hit = self._search_cache.get(norm_query)
+            if cache_hit:
+                return cache_hit
+            
+            # 确保索引已构建（首次或过期后会重建）
+            self._ensure_company_index()
+            
+            # 获取索引引用（已在锁内构建完成）
+            by_ticker = EdgarDataClient._by_ticker
+            by_title = EdgarDataClient._by_title
+            by_title_norm = EdgarDataClient._by_title_norm
+            all_keys = EdgarDataClient._all_keys
+            
+            # ✅ OPTIMIZATION 1: Ticker 优先匹配（遵循项目规范）
+            raw = company_name.strip().lower()
+            raw_compact = re.sub(r"[^a-z0-9]", "", raw)
+            is_ticker_like = len(raw_compact) <= 5 and len(raw_compact) >= 1
+            
+            if is_ticker_like and raw_compact in by_ticker:
+                result = by_ticker[raw_compact]
+                self._search_cache[norm_query] = result
+                return result
+            
+            # ✅ OPTIMIZATION 2: 别名映射（如 'google' -> 'alphabet inc'）
+            alias_target = EdgarDataClient._alias_map.get(norm_query)
+            if alias_target:
+                alias_norm = self._normalize_text(alias_target)
+                # 先尝试规范化标题
+                if alias_norm in by_title_norm:
+                    result = by_title_norm[alias_norm]
+                    self._search_cache[norm_query] = result
+                    return result
+                # 再尝试原始标题
+                alias_lower = alias_target.lower()
+                if alias_lower in by_title:
+                    result = by_title[alias_lower]
+                    self._search_cache[norm_query] = result
+                    return result
+                # 最后尝试 ticker（有些别名可能实际上是ticker）
+                alias_ticker = re.sub(r"[^a-z0-9]", "", alias_lower)
+                if alias_ticker in by_ticker:
+                    result = by_ticker[alias_ticker]
+                    self._search_cache[norm_query] = result
+                    return result
+            
+            # ✅ OPTIMIZATION 3: 精确匹配（原始标题）
+            title_lower = company_name.lower().strip()
+            if title_lower in by_title:
+                result = by_title[title_lower]
+                self._search_cache[norm_query] = result
+                return result
+            
+            # ✅ OPTIMIZATION 4: 精确匹配（规范化标题）
+            if norm_query in by_title_norm:
+                result = by_title_norm[norm_query]
+                self._search_cache[norm_query] = result
+                return result
+            
+            # ✅ OPTIMIZATION 5: 精确匹配（ticker，再次尝试原始输入）
+            if raw_compact in by_ticker:
+                result = by_ticker[raw_compact]
+                self._search_cache[norm_query] = result
+                return result
+            
+            # ✅ OPTIMIZATION 6: 部分包含匹配
+            partial_matches = []
+            for key in by_title_norm.keys():
+                if norm_query in key:
+                    partial_matches.append(key)
+            if not partial_matches:
+                for t in by_ticker.keys():
+                    if norm_query in t:
+                        partial_matches.append(t)
+            if partial_matches:
+                best_key = max(
+                    partial_matches,
+                    key=lambda k: difflib.SequenceMatcher(None, norm_query, k).ratio()
+                )
+                result = by_title_norm.get(best_key) or by_ticker.get(best_key)
+                if result:
+                    self._search_cache[norm_query] = result
+                    return result
+            
+            # ✅ OPTIMIZATION 7: 模糊匹配（difflib，用于拼写近似的情况）
+            close = difflib.get_close_matches(norm_query, all_keys, n=1, cutoff=0.78)
+            if close:
+                best = close[0]
+                result = by_title_norm.get(best) or by_ticker.get(best)
+                if result:
+                    self._search_cache[norm_query] = result
+                    return result
+            
+            # 未找到
+            return None
             
         except TimeoutError as e:
             print(f"Timeout searching company: {e}")