roy2012/stock_mcp

1# Utility functions, including the Baostock login context manager and logging setup2import baostock as bs3 4import os5import sys6import logging7from contextlib import contextmanager8from .data_source_interface import LoginError9 10# --- Logging Setup ---11def setup_logging(level=logging.INFO):12    """Configures basic logging for the application."""13    logging.basicConfig(14        level=level,15        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',16        datefmt='%Y-%m-%d %H:%M:%S'17    )18    # Optionally silence logs from dependencies if they are too verbose19    # logging.getLogger("mcp").setLevel(logging.WARNING)20 21# Get a logger instance for this module (optional, but good practice)22logger = logging.getLogger(__name__)23 24 25@contextmanager26def tushare_login_context():27    """Context manager to handle Tushare login and API initialization."""28    # Redirect stdout to suppress any potential messages29    original_stdout_fd = sys.stdout.fileno()30    saved_stdout_fd = os.dup(original_stdout_fd)31    devnull_fd = os.open(os.devnull, os.O_WRONLY)32 33    os.dup2(devnull_fd, original_stdout_fd)34    os.close(devnull_fd)35 36    import tushare as ts37    from dotenv import load_dotenv38    39    # 加载环境变量40    load_dotenv()41    token = os.getenv('TUSHARE_TOKEN')42    if not token:43        raise LoginError("请在 .env 文件中设置 TUSHARE_TOKEN")44 45    # 设置 token46    ts.set_token(token)47    # 初始化 pro 接口48    pro = ts.pro_api()49 50    # Restore stdout51    os.dup2(saved_stdout_fd, original_stdout_fd)52    os.close(saved_stdout_fd)53    try:54        yield pro  # 返回 pro 接口实例供使用55    except Exception as e:56        raise LoginError(f"Tushare login failed: {str(e)}")57 58# --- Baostock Context Manager ---59@contextmanager60def baostock_login_context():61    """Context manager to handle Baostock login and logout, suppressing stdout messages."""62    # Redirect stdout to suppress login/logout messages63    original_stdout_fd = sys.stdout.fileno()64    saved_stdout_fd = os.dup(original_stdout_fd)65    devnull_fd = os.open(os.devnull, os.O_WRONLY)66 67    os.dup2(devnull_fd, original_stdout_fd)68    os.close(devnull_fd)69 70    logger.debug("Attempting Baostock login...")71    lg = bs.login()72    logger.debug(f"Login result: code={lg.error_code}, msg={lg.error_msg}")73 74    # Restore stdout75    os.dup2(saved_stdout_fd, original_stdout_fd)76    os.close(saved_stdout_fd)77 78    if lg.error_code != '0':79        # Log error before raising80        logger.error(f"Baostock login failed: {lg.error_msg}")81        raise LoginError(f"Baostock login failed: {lg.error_msg}")82 83    logger.info("Baostock login successful.")84    try:85        yield  # API calls happen here86    finally:87        # Redirect stdout again for logout88        original_stdout_fd = sys.stdout.fileno()89        saved_stdout_fd = os.dup(original_stdout_fd)90        devnull_fd = os.open(os.devnull, os.O_WRONLY)91 92        os.dup2(devnull_fd, original_stdout_fd)93        os.close(devnull_fd)94 95        logger.debug("Attempting Baostock logout...")96        bs.logout()97        logger.debug("Logout completed.")98 99        # Restore stdout100        os.dup2(saved_stdout_fd, original_stdout_fd)101        os.close(saved_stdout_fd)102        logger.info("Baostock logout successful.")103 104# You can add other utility functions or classes here if needed105