yfinance_plus/yfinance_plus.py

#!/usr/bin/env python3
"""
Enhanced yfinance wrapper with improved rate limit handling and intelligent caching.
Uses dynamic proxy pattern to wrap all yfinance functionality automatically.
"""

import os
import time
import random
import threading
import pickle
import hashlib
from pathlib import Path
from typing import Dict, Optional, Any, Union, List
from dataclasses import dataclass
from functools import wraps
import logging
from datetime import datetime, timedelta

import yfinance as yf
from curl_cffi import requests
import pandas as pd


@dataclass
class RequestConfig:
    """Configuration for enhanced requests"""
    max_retries: int = 3
    base_delay: float = 1.0
    max_delay: float = 60.0
    jitter: bool = True
    user_agents: list = None
    enable_cache: bool = True
    cache_dir: str = None
    
    def __post_init__(self):
        if self.user_agents is None:
            self.user_agents = [
                'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36',
                'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36',
                'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36',
                'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.2.1 Safari/605.1.15',
                'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:121.0) Gecko/20100101 Firefox/121.0',
                'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36',
                'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0'
            ]
        if self.cache_dir is None:
            self.cache_dir = os.path.expanduser('~/.yfinance_cache')


class HistoricalDataCache:
    """Intelligent cache for historical data"""
    
    def __init__(self, cache_dir: str):
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self._lock = threading.Lock()
        
    def _get_cache_key(self, symbol: str, period: str, interval: str, start: str = None, end: str = None) -> str:
        """Generate cache key for historical data request"""
        key_data = f"{symbol}_{period}_{interval}_{start}_{end}"
        return hashlib.md5(key_data.encode()).hexdigest()
    
    def _get_cache_file(self, cache_key: str) -> Path:
        """Get cache file path"""
        return self.cache_dir / f"{cache_key}.pkl"
    
    def _normalize_period_to_days(self, period: str) -> int:
        """Convert period string to approximate days"""
        period_map = {
            '1d': 1, '2d': 2, '5d': 5, '1mo': 30, '3mo': 90, '6mo': 180,
            '1y': 365, '2y': 730, '5y': 1825, '10y': 3650, 'ytd': 365, 'max': 36500
        }
        return period_map.get(period.lower(), 30)
    
    def get(self, symbol: str, period: str, interval: str, start: str = None, end: str = None) -> Optional[pd.DataFrame]:
        """Get cached historical data if available and valid"""
        with self._lock:
            try:
                # First check exact match
                cache_key = self._get_cache_key(symbol, period, interval, start, end)
                cache_file = self._get_cache_file(cache_key)
                
                if cache_file.exists():
                    with open(cache_file, 'rb') as f:
                        cache_data = pickle.load(f)
                        
                    # Check if cache is still valid (not older than 1 hour for recent data)
                    cache_time = datetime.fromtimestamp(cache_file.stat().st_mtime)
                    if datetime.now() - cache_time < timedelta(hours=1):
                        return cache_data['data']
                
                # Look for longer period cache that might contain requested data
                requested_days = self._normalize_period_to_days(period)
                
                for cache_file in self.cache_dir.glob(f"*.pkl"):
                    try:
                        with open(cache_file, 'rb') as f:
                            cache_data = pickle.load(f)
                        
                        # Check if this cache contains data for the same symbol and interval
                        if (cache_data['symbol'] == symbol and 
                            cache_data['interval'] == interval and
                            cache_data['period_days'] >= requested_days):
                            
                            # Check cache validity
                            cache_time = datetime.fromtimestamp(cache_file.stat().st_mtime)
                            if datetime.now() - cache_time < timedelta(hours=1):
                                # Extract requested period from cached data
                                data = cache_data['data']
                                if len(data) >= requested_days:
                                    return data.tail(min(len(data), requested_days * 2))  # Buffer for weekends
                    except:
                        continue
                        
                return None
                
            except Exception as e:
                logging.warning(f"Cache read error: {e}")
                return None
    
    def store(self, symbol: str, period: str, interval: str, data: pd.DataFrame, start: str = None, end: str = None):
        """Store historical data in cache"""
        with self._lock:
            try:
                # Check if we already have a longer period cache that contains this data
                requested_days = self._normalize_period_to_days(period)
                
                # Look for existing caches that might already contain this data
                for cache_file in self.cache_dir.glob(f"*.pkl"):
                    try:
                        with open(cache_file, 'rb') as f:
                            existing_cache = pickle.load(f)
                        
                        # Check if existing cache already covers this request
                        if (existing_cache['symbol'] == symbol and 
                            existing_cache['interval'] == interval and
                            existing_cache['period_days'] >= requested_days):
                            
                            # Check if existing cache is recent enough
                            cache_time = datetime.fromtimestamp(cache_file.stat().st_mtime)
                            if datetime.now() - cache_time < timedelta(hours=1):
                                # We already have this data in a longer period cache
                                logging.info(f"Skipping cache store - data already exists in {cache_file.name}")
                                return
                    except:
                        continue
                
                # Store new cache if no suitable existing cache found
                cache_key = self._get_cache_key(symbol, period, interval, start, end)
                cache_file = self._get_cache_file(cache_key)
                
                cache_data = {
                    'symbol': symbol,
                    'period': period,
                    'period_days': self._normalize_period_to_days(period),
                    'interval': interval,
                    'start': start,
                    'end': end,
                    'data': data,
                    'cached_at': datetime.now()
                }
                
                with open(cache_file, 'wb') as f:
                    pickle.dump(cache_data, f)
                    
                # Clean up redundant caches (smaller periods for same symbol/interval)
                self._cleanup_redundant_caches(symbol, interval, requested_days, cache_key)
                    
            except Exception as e:
                logging.warning(f"Cache store error: {e}")
    
    def _cleanup_redundant_caches(self, symbol: str, interval: str, new_period_days: int, new_cache_key: str):
        """Clean up redundant caches that are covered by the new cache"""
        try:
            for cache_file in self.cache_dir.glob(f"*.pkl"):
                # Skip the new cache file
                if cache_file.name == f"{new_cache_key}.pkl":
                    continue
                    
                try:
                    with open(cache_file, 'rb') as f:
                        cache_data = pickle.load(f)
                    
                    # Remove smaller period caches for the same symbol/interval
                    if (cache_data['symbol'] == symbol and 
                        cache_data['interval'] == interval and
                        cache_data['period_days'] < new_period_days):
                        
                        cache_file.unlink()
                        logging.info(f"Removed redundant cache: {cache_file.name}")
                except:
                    continue
        except Exception as e:
            logging.warning(f"Cache cleanup error: {e}")
    
    def clear(self):
        """Clear all cached data"""
        with self._lock:
            for cache_file in self.cache_dir.glob("*.pkl"):
                try:
                    cache_file.unlink()
                except:
                    pass
    
    def get_cache_info(self) -> Dict[str, Any]:
        """Get cache statistics"""
        with self._lock:
            cache_files = list(self.cache_dir.glob("*.pkl"))
            total_size = sum(f.stat().st_size for f in cache_files)
            
            return {
                'cache_dir': str(self.cache_dir),
                'file_count': len(cache_files),
                'total_size_mb': total_size / (1024 * 1024),
                'files': [{
                    'name': f.name,
                    'size_kb': f.stat().st_size / 1024,
                    'modified': datetime.fromtimestamp(f.stat().st_mtime)
                } for f in cache_files]
            }


class EnhancedYFinance:
    """Enhanced yfinance wrapper with better rate limit handling and caching"""
    
    def __init__(self, config: RequestConfig = None):
        self.config = config or RequestConfig()
        self._session = None
        self._lock = threading.Lock()
        self._last_request_time = 0
        self._request_count = 0
        self._session_created_time = time.time()
        
        # Setup caching
        if self.config.enable_cache:
            self._cache = HistoricalDataCache(self.config.cache_dir)
        else:
            self._cache = None
        
        # Setup logging
        self.logger = logging.getLogger(__name__)
        if not self.logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            self.logger.addHandler(handler)
            self.logger.setLevel(logging.INFO)  # Show cache operations
            
        # Suppress yfinance HTTP error logs
        yf_logger = logging.getLogger('yfinance')
        yf_logger.setLevel(logging.CRITICAL)
    
    def _create_enhanced_session(self) -> requests.Session:
        """Create a session with browser-like headers and behavior"""
        session = requests.Session(impersonate="chrome")
        
        # Enhanced headers to mimic real browser behavior
        browser_headers = {
            'User-Agent': random.choice(self.config.user_agents),
            'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
            'Accept-Language': 'en-US,en;q=0.9,ko;q=0.8,ja;q=0.7',
            'Accept-Encoding': 'gzip, deflate, br',
            'Cache-Control': 'no-cache',
            'Pragma': 'no-cache',
            'Sec-Ch-Ua': '"Not_A Brand";v="8", "Chromium";v="120", "Google Chrome";v="120"',
            'Sec-Ch-Ua-Mobile': '?0',
            'Sec-Ch-Ua-Platform': '"macOS"',
            'Sec-Fetch-Dest': 'document',
            'Sec-Fetch-Mode': 'navigate',
            'Sec-Fetch-Site': 'cross-site',
            'Sec-Fetch-User': '?1',
            'Upgrade-Insecure-Requests': '1',
            'Connection': 'keep-alive',
            'DNT': '1',
            'Origin': 'https://finance.yahoo.com',
            'Referer': 'https://finance.yahoo.com/',
        }
        
        session.headers.update(browser_headers)
        
        # Add realistic Yahoo Finance cookies
        current_time = int(time.time())
        session.cookies.update({
            'A1': f'd=AQABBC{random.randint(100000, 999999)}YLoCIgEBwQJ7vgAB&S=AQAAAg',
            'A1S': f'd=AQABBC{random.randint(100000, 999999)}YLoCIgEBwQJ7vgAB&S=AQAAAg',
            'A3': f'd=AQABBC{random.randint(100000, 999999)}YLoCIgEBwQJ7vgAB&S=AQAAAg',
            'GUC': f'AQEBAQFm{random.randint(1000, 9999)}k0L',
            'B': f'c={random.randint(1000000, 9999999)}&b=3&s=4u',
            'cmp': f't={current_time}&j=0&u=1---',
            'EuConsent': 'CP-r9cAP-r9cAAOACKENAoEgAAAAAAAAACiQAAAAAAAA',
        })
        
        return session
    
    @property
    def session(self) -> requests.Session:
        """Get or create enhanced session"""
        with self._lock:
            if self._session is None or self._should_refresh_session():
                self._session = self._create_enhanced_session()
                self._session_created_time = time.time()
                self._request_count = 0
                self.logger.debug("Created new enhanced session")
            return self._session
    
    def _should_refresh_session(self) -> bool:
        """Check if session should be refreshed"""
        session_age = time.time() - self._session_created_time
        return (session_age > 600 or  # 10 minutes
                self._request_count > 50)  # 50 requests
    
    def _rate_limit_delay(self):
        """Apply intelligent rate limiting"""
        with self._lock:
            current_time = time.time()
            time_since_last = current_time - self._last_request_time
            
            # Minimum delay based on request frequency
            min_delay = 0.1 if self._request_count < 10 else 0.2
            
            if time_since_last < min_delay:
                delay = min_delay - time_since_last
                if self.config.jitter:
                    delay += random.uniform(0, delay * 0.5)
                time.sleep(delay)
            
            self._last_request_time = time.time()
            self._request_count += 1
    
    def _retry_with_backoff(self, func, *args, **kwargs):
        """Execute function with exponential backoff retry"""
        last_exception = None
        
        for attempt in range(self.config.max_retries + 1):
            try:
                self._rate_limit_delay()
                return func(*args, **kwargs)
            
            except Exception as e:
                last_exception = e
                error_str = str(e).lower()
                
                # Check if it's a rate limit or 401 error
                if ("rate limit" in error_str or "429" in error_str or 
                    "401" in error_str or "unauthorized" in error_str):
                    
                    if attempt < self.config.max_retries:
                        delay = min(
                            self.config.base_delay * (2 ** attempt),
                            self.config.max_delay
                        )
                        if self.config.jitter:
                            delay += random.uniform(0, delay * 0.3)
                        
                        if "401" in error_str:
                            self.logger.debug(f"401 error, refreshing session and retrying in {delay:.2f}s (attempt {attempt + 1})")
                        else:
                            self.logger.warning(f"Rate limit hit, retrying in {delay:.2f}s (attempt {attempt + 1})")
                        
                        time.sleep(delay)
                        
                        # Refresh session on auth/rate limit error
                        with self._lock:
                            self._session = None
                        continue
                else:
                    # Non-rate-limit error, re-raise immediately
                    raise e
        
        # All retries exhausted
        raise last_exception
    
    def get_ticker(self, symbol: str) -> 'EnhancedTicker':
        """Get enhanced ticker with improved rate limiting"""
        return EnhancedTicker(symbol, self)


class EnhancedTicker:
    """
    Enhanced ticker wrapper using dynamic proxy pattern.
    Automatically wraps ALL yfinance ticker functionality.
    """
    
    # Methods that should be cached
    _CACHEABLE_METHODS = {'history'}
    
    # Methods that should NOT be cached (real-time or frequently changing data)
    _NON_CACHEABLE_METHODS = {'news', 'get_news', 'fast_info'}
    
    def __init__(self, symbol: str, enhanced_yf: EnhancedYFinance):
        self.symbol = symbol.upper()
        self.enhanced_yf = enhanced_yf
        self._yf_ticker = None
    
    @property
    def yf_ticker(self):
        """Get yfinance ticker with enhanced session"""
        if self._yf_ticker is None:
            self._yf_ticker = yf.Ticker(self.symbol, session=self.enhanced_yf.session)
        return self._yf_ticker
    
    def __getattr__(self, name):
        """
        Dynamic proxy: automatically wrap any yfinance method or property
        """
        # Get the original attribute from yfinance ticker
        original_attr = getattr(self.yf_ticker, name)
        
        # If it's a method/function, wrap it with enhanced functionality
        if callable(original_attr):
            if name == 'history':
                # Special handling for history method with caching
                return self._enhanced_history
            elif name in self._NON_CACHEABLE_METHODS:
                # Methods that should not be cached
                return self._wrap_method_no_cache(original_attr)
            else:
                # All other methods get standard enhancement
                return self._wrap_method(original_attr)
        else:
            # For properties, wrap them with retry logic
            return self._wrap_property(name)
    
    def _enhanced_history(self, period="1mo", interval="1d", start=None, end=None, **kwargs):
        """Enhanced history method with caching"""
        # Check cache first for historical data (not real-time)
        if (self.enhanced_yf._cache and 
            not kwargs.get('prepost', False)):    # Don't cache extended hours data
            
            cached_data = self.enhanced_yf._cache.get(
                self.symbol, period, interval, 
                str(start) if start else None, 
                str(end) if end else None
            )
            
            if cached_data is not None and not cached_data.empty:
                self.enhanced_yf.logger.info(f"Using cached data for {self.symbol}")
                return cached_data
        
        # Fetch fresh data
        def _get_history():
            return self.yf_ticker.history(
                period=period, interval=interval, 
                start=start, end=end, **kwargs
            )
        
        data = self.enhanced_yf._retry_with_backoff(_get_history)
        
        # Cache the data if successful and caching is enabled
        if (self.enhanced_yf._cache and 
            data is not None and not data.empty and
            not kwargs.get('prepost', False)):
            
            self.enhanced_yf._cache.store(
                self.symbol, period, interval, data,
                str(start) if start else None,
                str(end) if end else None
            )
        
        return data
    
    def _wrap_method(self, method):
        """Wrap a method with enhanced error handling and rate limiting"""
        @wraps(method)
        def wrapped(*args, **kwargs):
            def _call_method():
                return method(*args, **kwargs)
            return self.enhanced_yf._retry_with_backoff(_call_method)
        return wrapped
    
    def _wrap_method_no_cache(self, method):
        """Wrap a method with enhanced error handling but no caching"""
        @wraps(method)
        def wrapped(*args, **kwargs):
            def _call_method():
                return method(*args, **kwargs)
            return self.enhanced_yf._retry_with_backoff(_call_method)
        return wrapped
    
    def _wrap_property(self, name):
        """Wrap a property with enhanced error handling"""
        def _get_property():
            return getattr(self.yf_ticker, name)
        return self.enhanced_yf._retry_with_backoff(_get_property)
    
    # Additional utility methods
    def get_cache_info(self) -> Dict[str, Any]:
        """Get cache information for this ticker"""
        if self.enhanced_yf._cache:
            return self.enhanced_yf._cache.get_cache_info()
        return {'cache_enabled': False}
    
    def clear_cache(self):
        """Clear cache for this ticker"""
        if self.enhanced_yf._cache:
            self.enhanced_yf._cache.clear()


def download(tickers: Union[str, list], 
             period: str = "1mo", 
             interval: str = "1d",
             **kwargs) -> pd.DataFrame:
    """Enhanced download function with better rate limiting"""
    
    config = RequestConfig(
        max_retries=5,
        base_delay=1.0,
        max_delay=120.0
    )
    
    enhanced_yf = EnhancedYFinance(config)
    
    def _download():
        return yf.download(
            tickers=tickers,
            period=period,
            interval=interval,
            session=enhanced_yf.session,
            **kwargs
        )
    
    return enhanced_yf._retry_with_backoff(_download)


# Global configuration
_global_config = RequestConfig()

def set_config(**kwargs):
    """Set global configuration for enhanced yfinance"""
    global _global_config
    for key, value in kwargs.items():
        if hasattr(_global_config, key):
            setattr(_global_config, key, value)
        else:
            raise ValueError(f"Unknown configuration option: {key}")

def get_config() -> RequestConfig:
    """Get current global configuration"""
    return _global_config

# Convenience functions with full yfinance API compatibility
def Ticker(symbol: str, session=None, proxy=None) -> EnhancedTicker:
    """Create enhanced ticker (fully compatible with yfinance.Ticker)"""
    config = RequestConfig(
        max_retries=_global_config.max_retries,
        base_delay=_global_config.base_delay,
        max_delay=_global_config.max_delay,
        jitter=_global_config.jitter,
        user_agents=_global_config.user_agents,
        enable_cache=_global_config.enable_cache,
        cache_dir=_global_config.cache_dir
    )
    
    enhanced_yf = EnhancedYFinance(config)
    
    # Override session if provided (for compatibility)
    if session is not None:
        enhanced_yf._session = session
    
    return enhanced_yf.get_ticker(symbol)

class EnhancedTickers:
    """Enhanced Tickers class with full yfinance compatibility"""
    
    def __init__(self, tickers, session=None, proxy=None):
        # Parse tickers like original yfinance
        if isinstance(tickers, str):
            tickers = tickers.replace(',', ' ').split()
        elif not isinstance(tickers, list):
            tickers = list(tickers)
            
        self.symbols = [ticker.upper() for ticker in tickers]
        self.tickers = {symbol: Ticker(symbol, session, proxy) for symbol in self.symbols}
        
        # Store original yfinance Tickers for method delegation
        self._yf_tickers = yf.Tickers(self.symbols, session=session)
    
    def __repr__(self):
        return f"yfinance_enhanced.Tickers object <{','.join(self.symbols)}>"
    
    def __getattr__(self, name):
        """Delegate any missing methods to original yfinance Tickers"""
        return getattr(self._yf_tickers, name)
    
    def download(self, *args, **kwargs):
        """Enhanced download for multiple tickers"""
        return download(self.symbols, *args, **kwargs)

# Additional yfinance compatibility functions  
def Tickers(tickers, session=None, proxy=None) -> EnhancedTickers:
    """Create enhanced tickers (fully compatible with yfinance.Tickers)"""
    return EnhancedTickers(tickers, session, proxy)

# Cache management functions
def clear_cache():
    """Clear all cached data"""
    cache = HistoricalDataCache(_global_config.cache_dir)
    cache.clear()
    
def get_cache_info():
    """Get cache information"""
    cache = HistoricalDataCache(_global_config.cache_dir)
    return cache.get_cache_info()

# Make the module work exactly like yfinance
from yfinance import __version__

# Export main classes and functions for full compatibility
__all__ = [
    'Ticker', 'Tickers', 'download', 'set_config', 'get_config',
    'clear_cache', 'get_cache_info', 'RequestConfig', 
    'EnhancedTicker', 'EnhancedTickers', 'EnhancedYFinance', 'HistoricalDataCache'
]


# Example usage
if __name__ == "__main__":
    # Setup logging
    logging.basicConfig(level=logging.INFO)
    
    # Create enhanced ticker
    ticker = Ticker("AAPL")
    
    try:
        # Test all yfinance functionality automatically
        print("Testing enhanced yfinance with automatic wrapping...")
        
        # All these should work automatically via __getattr__
        print(f"Info: {ticker.info.get('longName', 'N/A')}")
        print(f"History: {len(ticker.history(period='5d'))} rows")
        print(f"Dividends: {len(ticker.dividends)} records")
        print(f"Splits: {len(ticker.splits)} records")
        print(f"Major holders: {ticker.major_holders.shape if hasattr(ticker.major_holders, 'shape') else 'Available'}")
        print(f"Recommendations: {ticker.recommendations.shape if hasattr(ticker.recommendations, 'shape') else 'Available'}")
        print(f"News: {len(ticker.news) if ticker.news else 0} articles")
        
        print("✅ All yfinance functionality works automatically!")
        
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()