spb/uqo-research-db
Public
1# ============================================================2# Project: EDGAR Research Database3# Description: SEC EDGAR Research Infrastructure4#5# Authors:6# - Manel Kammoun — manel.kammoun@uqo.ca7# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca8# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca9#10# Institution:11# Université du Québec en Outaouais (UQO)12#13# License:14# Research / Academic Use15# ============================================================1617from __future__ import annotations1819import json20import os21import time22from datetime import datetime, timezone23from pathlib import Path2425import httpx26from tqdm import tqdm2728from src.acquisition.sec_client import SECClient29from src.utils.constants import DATA_DIRS30from src.utils.logging import get_logger3132logger = get_logger("acquisition.downloader")3334MIN_REQUEST_INTERVAL = 0.135RETRY_BACKOFF_BASE = 2.036MAX_RETRIES = 5373839class FilingDownloader:4041 def __init__(self, output_dir: Path | None = None) -> None:42 self._output_dir = output_dir or DATA_DIRS["raw"]43 self._metadata_dir = DATA_DIRS["metadata"]44 self._download_log_path = self._metadata_dir / "download_log.json"45 self._user_agent = os.environ.get(46 "SEC_USER_AGENT", "Research Agent research@example.com"47 )48 self._last_request_time: float = 0.049 self._downloaded: set[str] = set()50 self._client = httpx.Client(51 timeout=30,52 headers={"User-Agent": self._user_agent},53 follow_redirects=True,54 )55 self._load_download_log()5657 def close(self) -> None:58 self._client.close()5960 def __enter__(self):61 return self6263 def __exit__(self, *args):64 self.close()6566 def download_filing(67 self,68 cik: str,69 accession: str,70 primary_doc: str,71 form_type: str = "10-K",72 filing_url: str | None = None,73 ) -> Path:74 padded_cik = cik.zfill(10)75 accession_clean = accession.replace("-", "")76 filing_date_year = ""7778 if filing_url:79 url = filing_url80 else:81 raw_cik = cik.lstrip("0") or "0"82 url = (83 f"https://www.sec.gov/Archives/edgar/data/"84 f"{raw_cik}/{accession_clean}/{primary_doc}"85 )8687 save_dir = (88 self._output_dir89 / padded_cik90 / form_type91 / accession_clean92 )93 save_dir.mkdir(parents=True, exist_ok=True)94 save_path = save_dir / primary_doc9596 if save_path.exists():97 logger.debug("File already exists: %s", save_path)98 return save_path99100 content = self._download_with_retry(url)101 if content is None:102 raise RuntimeError(f"Failed to download {url} after {MAX_RETRIES} retries")103104 save_path.write_bytes(content)105 logger.debug("Downloaded %s -> %s", url, save_path)106 return save_path107108 def download_company_filings(109 self,110 cik: str,111 ticker: str,112 form_type: str = "10-K",113 start_date: str = "2010-01-01",114 sec_client: SECClient | None = None,115 ) -> list[dict]:116 own_client = sec_client is None117 client = sec_client or SECClient()118 try:119 filings = client.get_company_filings(120 cik=cik, form_type=form_type, start_date=start_date121 )122 finally:123 if own_client:124 client.close()125126 results: list[dict] = []127 for filing in filings:128 accession = filing["accession_number"]129 download_key = f"{cik}|{accession}|{filing['primary_document']}"130131 if download_key in self._downloaded:132 logger.debug("Skipping already-downloaded: %s", download_key)133 filing["local_path"] = str(134 self._output_dir135 / cik.zfill(10)136 / form_type137 / accession.replace("-", "")138 / filing["primary_document"]139 )140 filing["ticker"] = ticker141 results.append(filing)142 continue143144 try:145 local_path = self.download_filing(146 cik=cik,147 accession=accession,148 primary_doc=filing["primary_document"],149 form_type=form_type,150 filing_url=filing.get("filing_url"),151 )152 filing["local_path"] = str(local_path)153 filing["ticker"] = ticker154 results.append(filing)155 self._record_download(download_key, filing)156 except Exception:157 logger.error(158 "Failed to download filing %s for %s (%s)",159 accession, ticker, cik,160 exc_info=True,161 )162163 logger.info(164 "Downloaded %d/%d %s filings for %s (CIK %s)",165 len(results), len(filings), form_type, ticker, cik,166 )167 return results168169 def download_universe(170 self,171 tickers: list[str],172 form_types: list[str],173 start_date: str = "2010-01-01",174 resume: bool = True,175 cik_mapper=None,176 ) -> None:177 if cik_mapper is None:178 from src.acquisition.cik_mapper import CIKMapper179 cik_mapper = CIKMapper()180181 if not resume:182 self._downloaded.clear()183184 total_tasks = len(tickers) * len(form_types)185 progress = tqdm(total=total_tasks, desc="Downloading universe", unit="task")186 completed = 0187 failed = 0188189 sec_client = SECClient()190 try:191 for ticker in tickers:192 cik = cik_mapper.get_cik(ticker)193 if not cik:194 logger.warning("No CIK found for ticker %s, skipping", ticker)195 progress.update(len(form_types))196 failed += len(form_types)197 continue198199 for form_type in form_types:200 progress.set_postfix_str(f"{ticker} {form_type}")201 try:202 self.download_company_filings(203 cik=cik,204 ticker=ticker,205 form_type=form_type,206 start_date=start_date,207 sec_client=sec_client,208 )209 completed += 1210 except Exception:211 logger.error(212 "Failed universe download for %s %s",213 ticker, form_type,214 exc_info=True,215 )216 failed += 1217 progress.update(1)218 finally:219 sec_client.close()220221 progress.close()222 logger.info(223 "Universe download complete: %d succeeded, %d failed out of %d tasks",224 completed, failed, total_tasks,225 )226227 def _download_with_retry(self, url: str) -> bytes | None:228 for attempt in range(MAX_RETRIES):229 self._rate_limit()230 try:231 response = self._client.get(url)232 if response.status_code == 404:233 logger.warning("Filing not found (404): %s", url)234 return None235 response.raise_for_status()236 return response.content237 except httpx.HTTPStatusError as exc:238 wait = RETRY_BACKOFF_BASE ** attempt239 logger.warning(240 "HTTP %d downloading %s (attempt %d/%d), retrying in %.1fs",241 exc.response.status_code, url, attempt + 1, MAX_RETRIES, wait,242 )243 time.sleep(wait)244 except httpx.RequestError as exc:245 wait = RETRY_BACKOFF_BASE ** attempt246 logger.warning(247 "Request error downloading %s: %s (attempt %d/%d), retrying in %.1fs",248 url, exc, attempt + 1, MAX_RETRIES, wait,249 )250 time.sleep(wait)251252 logger.error("All %d retries exhausted for %s", MAX_RETRIES, url)253 return None254255 def _rate_limit(self) -> None:256 elapsed = time.monotonic() - self._last_request_time257 if elapsed < MIN_REQUEST_INTERVAL:258 time.sleep(MIN_REQUEST_INTERVAL - elapsed)259 self._last_request_time = time.monotonic()260261 def _load_download_log(self) -> None:262 if not self._download_log_path.exists():263 return264 try:265 with open(self._download_log_path) as f:266 for line in f:267 line = line.strip()268 if not line:269 continue270 entry = json.loads(line)271 key = entry.get("key", "")272 if key:273 self._downloaded.add(key)274 logger.info("Loaded %d entries from download log", len(self._downloaded))275 except Exception:276 logger.warning("Could not load download log, starting fresh", exc_info=True)277278 def _record_download(self, key: str, filing: dict) -> None:279 self._downloaded.add(key)280 self._metadata_dir.mkdir(parents=True, exist_ok=True)281 entry = {282 "key": key,283 "timestamp": datetime.now(timezone.utc).isoformat(),284 "ticker": filing.get("ticker", ""),285 "form_type": filing.get("form_type", ""),286 "accession_number": filing.get("accession_number", ""),287 "filing_date": filing.get("filing_date", ""),288 "local_path": filing.get("local_path", ""),289 }290 with open(self._download_log_path, "a") as f:291 f.write(json.dumps(entry) + "\n")292