scratch/scrape_all_fields.py

import asyncio
import json
import logging
from pathlib import Path
import os
import sys
from typing import Dict, Any, List

# Setup logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
logger = logging.getLogger("Scraper")

# Setup system paths
sys.path.append(str(Path(__file__).resolve().parent.parent))
import config
from core.wq_client import AsyncBrainClient, _parse_retry_after

async def scrape_fields():
    # Active regions, universes, and delays to capture all 100k+ fields
    regions = ["USA", "GLB", "EUR", "ASI", "CHN", "AMR", "IND", "MEA"]
    universes = ["TOP3000", "TOP2000", "TOP1000", "TOP500", "TOP200", "TOPSP500", "GLB3000", "GLB2000"]
    delays = [0, 1]

    # Load existing cache if any to build on top of it
    cache_file = config.DATA_DIR / "data_fields_cache.json"
    all_categorized_fields = {}
    if cache_file.exists():
        try:
            with open(cache_file, "r") as f:
                all_categorized_fields = json.load(f)
            total = sum(len(v) for v in all_categorized_fields.values())
            logger.info(f"Loaded existing cache containing {total} fields.")
        except Exception:
            all_categorized_fields = {}

    async with AsyncBrainClient() as client:
        # Authenticate first
        authenticated = await client.authenticate()
        if not authenticated:
            logger.error("Authentication with WorldQuant BRAIN failed. Please ensure credentials are set.")
            return

        session = await client.get_session()
        
        # Track unique field IDs
        seen_ids = set()
        for cat, fields in all_categorized_fields.items():
            for f in fields:
                seen_ids.add(f["id"])

        logger.info("Starting recursive pagination across all categories, regions, and universes...")
        
        # Paginate through combinations
        for region in regions:
            for universe in universes:
                for delay in delays:
                    logger.info(f"Scraping Region: {region}, Universe: {universe}, Delay: {delay}...")
                    offset = 0
                    limit = 50
                    consecutive_empty = 0
                    
                    while True:
                        url = f"{client.base_url}/data-fields"
                        params = {
                            "instrumentType": "EQUITY",
                            "region": region,
                            "delay": delay,
                            "universe": universe,
                            "limit": limit,
                            "offset": offset
                        }
                        
                        try:
                            async with session.get(url, params=params) as response:
                                if response.status == 401:
                                    logger.info("Session expired. Authenticating...")
                                    await client.authenticate()
                                    continue
                                    
                                if response.status == 429 or "Retry-After" in response.headers:
                                    retry_after = _parse_retry_after(response.headers, 15)
                                    logger.warning(f"Rate limited. Waiting {retry_after} seconds...")
                                    await asyncio.sleep(retry_after)
                                    continue
                                    
                                if response.status != 200:
                                    body = await response.text()
                                    logger.error(f"HTTP {response.status}: {body[:200]}")
                                    break
                                    
                                data = await response.json()
                                results = data.get("results", [])
                                if not results:
                                    consecutive_empty += 1
                                    if consecutive_empty >= 2:
                                        break
                                    offset += limit
                                    continue
                                    
                                consecutive_empty = 0
                                new_added = 0
                                
                                for field in results:
                                    fid = field.get("id")
                                    if not fid or fid in seen_ids:
                                        continue
                                        
                                    desc = field.get("description", "")
                                    category = field.get("category", {})
                                    cat_name = category.get("name", "Other") if isinstance(category, dict) else "Other"
                                    field_type = field.get("type", "MATRIX")
                                    
                                    if cat_name not in all_categorized_fields:
                                        all_categorized_fields[cat_name] = []
                                        
                                    all_categorized_fields[cat_name].append({
                                        "id": fid,
                                        "desc": desc,
                                        "type": field_type
                                    })
                                    seen_ids.add(fid)
                                    new_added += 1
                                    
                                total_now = sum(len(v) for v in all_categorized_fields.values())
                                logger.info(f"[{region}/{universe}] Fetched {len(results)} fields (offset {offset}). Added {new_added} new. Total Cache Size: {total_now}")
                                
                                if len(results) < limit:
                                    break
                                    
                                offset += limit
                                await asyncio.sleep(0.5)  # Politeness interval
                                
                        except Exception as e:
                            logger.error(f"Network error: {e}")
                            await asyncio.sleep(5)
                            break
                            
                    # Save progress after finishing a universe
                    if seen_ids:
                        with open(cache_file, "w") as f:
                            json.dump(all_categorized_fields, f, indent=2)

    total_scraped = sum(len(v) for v in all_categorized_fields.values())
    logger.info(f"Fields scrape complete! Total fields saved to cache: {total_scraped}")

if __name__ == "__main__":
    asyncio.run(scrape_fields())