news_data.py

"""
新闻公告数据获取模块
包含新闻和公告相关的数据获取方法
"""

import pandas as pd
from typing import List, Optional, Dict, Any
from concurrent.futures import ThreadPoolExecutor, as_completed
import time

# 设置环境变量，确保Python使用UTF-8编码处理所有I/O
import os
os.environ['PYTHONIOENCODING'] = 'utf-8'

# 使用共享日志模块
try:
    from .shared_logger import get_logger
except ImportError:
    from shared_logger import get_logger

# 设置日志
logger = get_logger(__name__)


class NewsDataMixin:
    """新闻公告数据获取混入类"""
    
    def news(self, src: str = '', start_date: str = '', end_date: str = '',
                fields: List[str] = None, use_cache: bool = True) -> pd.DataFrame:
        """
        获取新闻快讯（带智能缓存）
        
        Args:
            src: 新闻来源（可选）
            start_date: 开始日期（格式：20240101）
            end_date: 结束日期（格式：20241231）
            fields: 字段列表，可选
            use_cache: 是否使用缓存
            
        Returns:
            包含新闻快讯的DataFrame
        """
        # 生成缓存键
        symbol = f"news_{src}"
        cache_key = self._get_cache_key(
            data_type='news',
            symbol=symbol,
            start_date=start_date,
            end_date=end_date,
            src=src,
            fields=fields
        )
        
        # 检查缓存
        if use_cache:
            cached_data = self._get_from_cache(cache_key)
            if cached_data is not None:
                return cached_data
        
        # 默认字段
        if not fields:
            fields = ['datetime', 'content', 'title', 'channels', 'score']
        
        # 从Tushare获取数据
        if self.pro:
            try:
                df = self.pro.news(
                    src=src,
                    start_date=start_date,
                    end_date=end_date,
                    fields=','.join(fields)
                )
                
                if df.empty:
                    logger.error(f"❌ Tushare API返回空数据: news")
                    return pd.DataFrame()
                
                # 保存到缓存
                if use_cache:
                    self._save_to_cache(cache_key, df)
                
                return df
                
            except Exception as e:
                logger.error(f"获取新闻快讯失败: {e}")
                return pd.DataFrame()
        else:
            logger.error("❌ 未初始化Tushare Pro API")
            return pd.DataFrame()
    
    def notice(self, ann_id: str = '', ts_code: str = '', start_date: str = '',
                 end_date: str = '', fields: List[str] = None,
                 use_cache: bool = True) -> pd.DataFrame:
        """
        获取公告（带智能缓存）
        
        Args:
            ann_id: 公告ID
            ts_code: 股票代码（格式：000001.SZ）
            start_date: 开始日期（格式：20240101）
            end_date: 结束日期（格式：20241231）
            fields: 字段列表，可选
            use_cache: 是否使用缓存
            
        Returns:
            包含公告的DataFrame
        """
        # 生成缓存键
        symbol = f"notice_{ts_code}"
        cache_key = self._get_cache_key(
            data_type='notice',
            symbol=symbol,
            start_date=start_date,
            end_date=end_date,
            ann_id=ann_id,
            ts_code=ts_code,
            fields=fields
        )
        
        # 检查缓存
        if use_cache:
            cached_data = self._get_from_cache(cache_key)
            if cached_data is not None:
                return cached_data
        
        # 默认字段
        if not fields:
            fields = ['ann_id', 'ts_code', 'ann_date', 'title', 'ann_type', 
                     'ann_content', 'market', 'src', 'url']
        
        # 从Tushare获取数据
        if self.pro:
            try:
                df = self.pro.notice(
                    ann_id=ann_id,
                    ts_code=ts_code,
                    start_date=start_date,
                    end_date=end_date,
                    fields=','.join(fields)
                )
                
                if df.empty:
                    logger.error(f"❌ Tushare API返回空数据: notice")
                    return pd.DataFrame()
                
                # 保存到缓存
                if use_cache:
                    self._save_to_cache(cache_key, df)
                
                return df
                
            except Exception as e:
                logger.error(f"获取公告失败: {e}")
                return pd.DataFrame()
        else:
            logger.error("❌ 未初始化Tushare Pro API")
            return pd.DataFrame()
    
    def major_news(self, src: str = '', start_date: str = '', end_date: str = '',
                      fields: List[str] = None, use_cache: bool = True) -> pd.DataFrame:
        """
        获取要闻（带智能缓存）
        
        Args:
            src: 新闻来源（可选）
            start_date: 开始日期（格式：20240101）
            end_date: 结束日期（格式：20241231）
            fields: 字段列表，可选
            use_cache: 是否使用缓存
            
        Returns:
            包含要闻的DataFrame
        """
        # 生成缓存键
        symbol = f"major_news_{src}"
        cache_key = self._get_cache_key(
            data_type='major_news',
            symbol=symbol,
            start_date=start_date,
            end_date=end_date,
            src=src,
            fields=fields
        )
        
        # 检查缓存
        if use_cache:
            cached_data = self._get_from_cache(cache_key)
            if cached_data is not None:
                return cached_data
        
        # 默认字段
        if not fields:
            fields = ['datetime', 'content', 'title', 'channels', 'score']
        
        # 从Tushare获取数据
        if self.pro:
            try:
                df = self.pro.major_news(
                    src=src,
                    start_date=start_date,
                    end_date=end_date,
                    fields=','.join(fields)
                )
                
                if df.empty:
                    logger.error(f"❌ Tushare API返回空数据: major_news")
                    return pd.DataFrame()
                
                # 保存到缓存
                if use_cache:
                    self._save_to_cache(cache_key, df)
                
                return df
                
            except Exception as e:
                logger.error(f"获取要闻失败: {e}")
                return pd.DataFrame()
        else:
            logger.error("❌ 未初始化Tushare Pro API")
            return pd.DataFrame()
    
    def notice_batch(self, ts_codes: List[str], start_date: str = '', 
                        end_date: str = '', max_workers: int = 5) -> pd.DataFrame:
        """
        批量获取公告（并行处理）
        
        Args:
            ts_codes: 股票代码列表
            start_date: 开始日期（格式：20240101）
            end_date: 结束日期（格式：20241231）
            max_workers: 最大并行工作线程数
            
        Returns:
            包含所有公告的DataFrame
        """
        all_data = []
        
        def get_single_notice(ts_code):
            return self.notice(
                ts_code=ts_code,
                start_date=start_date,
                end_date=end_date
            )
        
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = {executor.submit(get_single_notice, ts_code): ts_code 
                      for ts_code in ts_codes}
            
            for future in as_completed(futures):
                ts_code = futures[future]
                try:
                    df = future.result()
                    if not df.empty:
                        all_data.append(df)
                except Exception as e:
                    logger.error(f"获取股票{ts_code}公告失败: {e}")
        
        if all_data:
            return pd.concat(all_data, ignore_index=True)
        else:
            return pd.DataFrame()