python_market_data_service/scripts/add_adapter_methods.py

#!/usr/bin/env python3
"""Add new data fetch methods to amazingdata_adapter.py using English comments"""

new_methods = '''

    # ==================== New Split Table Data Fetch Methods ====================
    
    async def fetch_kline_base(
        self,
        symbol: str,
        start: str,
        end: str,
        freq: str
    ) -> List[Dict[str, Any]]:
        """Fetch K-line base data (OHLCV)
        
        Corresponding tables: stock_klines_1d_base, stock_klines_1m_base, etc.
        
        Returns:
            List[Dict] containing fields:
            - symbol: Symbol code
            - ts: Timestamp
            - trade_date: Trade date
            - open/high/low/close: Price data
            - volume: Trading volume
            - amount: Trading amount
            - adj_factor: Adjustment factor
        """
        print(f"[amazingdata_adapter fetch_kline_base]Fetching {symbol} {freq} base data...")
        self._check_login()
        
        period_map = {
            "1m": self._ad.constant.Period.min1,
            "5m": self._ad.constant.Period.min5,
            "15m": self._ad.constant.Period.min15,
            "30m": self._ad.constant.Period.min30,
            "60m": self._ad.constant.Period.min60,
            "1d": self._ad.constant.Period.day,
            "1w": self._ad.constant.Period.week,
            "1month": self._ad.constant.Period.month,
        }
        period_value = period_map.get(freq, self._ad.constant.Period.day).value
        
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            None,
            lambda: self._fetch_kline_base_sync(symbol, start, end, period_value, freq)
        )
    
    def _fetch_kline_base_sync(
        self,
        symbol: str,
        start_date: str,
        end_date: str,
        period_value: int,
        freq: str
    ) -> List[Dict[str, Any]]:
        """Sync method to fetch K-line base data"""
        codes = [symbol]
        start_int = self._format_date(start_date)
        end_int = self._format_date(end_date)
        
        kline_dict = self._market_data.query_kline(
            code_list=codes,
            begin_date=start_int,
            end_date=end_int,
            period=period_value
        )
        
        if symbol not in kline_dict:
            info(f"No kline data found for {symbol}")
            return []
        
        df = kline_dict[symbol]
        results = []
        
        for _, row in df.iterrows():
            kline_time = row.get('kline_time')
            if pd.isna(kline_time) or kline_time is None:
                continue
            
            if isinstance(kline_time, pd.Timestamp):
                ts = int(kline_time.timestamp())
                trade_date = kline_time.strftime('%Y-%m-%d')
            else:
                date_str = str(int(kline_time))
                if len(date_str) != 8:
                    continue
                dt = datetime.strptime(date_str, "%Y%m%d")
                ts = int(dt.timestamp())
                trade_date = dt.strftime('%Y-%m-%d')
            
            results.append({
                "symbol": symbol,
                "ts": ts,
                "trade_date": trade_date,
                "open": float(row.get('open', 0)),
                "high": float(row.get('high', 0)),
                "low": float(row.get('low', 0)),
                "close": float(row.get('close', 0)),
                "volume": int(row.get('volume', 0)),
                "amount": float(row.get('amount', 0)),
                "adj_factor": float(row.get('adj_factor', 1.0)) if 'adj_factor' in df.columns else 1.0,
            })
        
        info(f"Fetched {len(results)} base kline records for {symbol}")
        return results
    
    async def fetch_kline_quote(
        self,
        symbol: str,
        start: str,
        end: str
    ) -> List[Dict[str, Any]]:
        """Fetch daily quote indicator data (calculated)
        
        Corresponding table: stock_klines_1d_quote
        
        Returns:
            List[Dict] containing fields:
            - change_pct: Price change percentage
            - change_Nd_pct: N-day price change
            - ma_N: Moving averages
            - macd_dif/dea/bar: MACD indicators
            - bias_N: Bias ratios
            - is_limit_up/down: Limit up/down status
            - is_st: ST status
        """
        print(f"[amazingdata_adapter fetch_kline_quote]Calculating {symbol} quote indicators...")
        self._check_login()
        
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            None,
            lambda: self._fetch_kline_quote_sync(symbol, start, end)
        )
    
    def _fetch_kline_quote_sync(
        self,
        symbol: str,
        start_date: str,
        end_date: str
    ) -> List[Dict[str, Any]]:
        """Sync method to calculate quote indicators"""
        import numpy as np
        
        start_dt = datetime.strptime(start_date, "%Y%m%d")
        extended_start = datetime(start_dt.year - 1, start_dt.month, start_dt.day)
        extended_start_str = extended_start.strftime("%Y%m%d")
        
        codes = [symbol]
        start_int = self._format_date(extended_start_str)
        end_int = self._format_date(end_date)
        
        kline_dict = self._market_data.query_kline(
            code_list=codes,
            begin_date=start_int,
            end_date=end_int,
            period=self._ad.constant.Period.day.value
        )
        
        if symbol not in kline_dict:
            return []
        
        df = kline_dict[symbol].copy()
        df = df.sort_values('kline_time')
        
        try:
            code_info_df = self._base_data.get_code_info(security_type=SecurityType.STOCK_A.value)
            if symbol in code_info_df.index:
                high_limited = float(code_info_df.loc[symbol, 'high_limited']) if 'high_limited' in code_info_df.columns else None
                low_limited = float(code_info_df.loc[symbol, 'low_limited']) if 'low_limited' in code_info_df.columns else None
            else:
                high_limited = low_limited = None
        except:
            high_limited = low_limited = None
        
        results = []
        closes = df['close'].values
        
        for i, (_, row) in enumerate(df.iterrows()):
            kline_time = row.get('kline_time')
            if pd.isna(kline_time):
                continue
            
            if isinstance(kline_time, pd.Timestamp):
                trade_date = kline_time.strftime('%Y-%m-%d')
            else:
                date_str = str(int(kline_time))
                if len(date_str) != 8:
                    continue
                dt = datetime.strptime(date_str, "%Y%m%d")
                trade_date = dt.strftime('%Y-%m-%d')
            
            close = float(row.get('close', 0))
            
            change_pct = None
            if i > 0:
                prev_close = closes[i-1]
                if prev_close > 0:
                    change_pct = round((close - prev_close) / prev_close * 100, 4)
            
            def calc_n_day_change(n):
                if i >= n and closes[i-n] > 0:
                    return round((close - closes[i-n]) / closes[i-n] * 100, 4)
                return None
            
            def calc_ma(n):
                if i >= n - 1:
                    return round(np.mean(closes[i-n+1:i+1]), 4)
                return None
            
            def calc_macd():
                if i < 33:
                    return None, None, None
                ema12 = pd.Series(closes[:i+1]).ewm(span=12).mean().iloc[-1]
                ema26 = pd.Series(closes[:i+1]).ewm(span=26).mean().iloc[-1]
                dif = ema12 - ema26
                dea = pd.Series([ema12 - ema26 for _ in range(i+1)]).ewm(span=9).mean().iloc[-1]
                bar = (dif - dea) * 2
                return round(dif, 6), round(dea, 6), round(bar, 6)
            
            def calc_bias(n):
                ma = calc_ma(n)
                if ma and ma > 0:
                    return round((close - ma) / ma * 100, 4)
                return None
            
            is_limit_up = False
            is_limit_down = False
            if high_limited and low_limited and close > 0:
                is_limit_up = close >= high_limited * 0.995
                is_limit_down = close <= low_limited * 1.005
            
            macd_dif, macd_dea, macd_bar = calc_macd()
            
            if trade_date.replace('-', '') >= start_date:
                results.append({
                    "symbol": symbol,
                    "trade_date": trade_date,
                    "change_pct": change_pct,
                    "change_5d_pct": calc_n_day_change(5),
                    "change_10d_pct": calc_n_day_change(10),
                    "change_20d_pct": calc_n_day_change(20),
                    "change_30d_pct": calc_n_day_change(30),
                    "change_60d_pct": calc_n_day_change(60),
                    "macd_dif": macd_dif,
                    "macd_dea": macd_dea,
                    "macd_bar": macd_bar,
                    "bias_5": calc_bias(5),
                    "bias_10": calc_bias(10),
                    "bias_20": calc_bias(20),
                    "is_limit_up": is_limit_up,
                    "is_limit_down": is_limit_down,
                    "limit_up_price": round(high_limited, 4) if high_limited else None,
                    "limit_down_price": round(low_limited, 4) if low_limited else None,
                    "is_st": None,
                    "ma_5": calc_ma(5),
                    "ma_10": calc_ma(10),
                    "ma_20": calc_ma(20),
                    "ma_30": calc_ma(30),
                    "ma_60": calc_ma(60),
                    "ma_120": calc_ma(120),
                    "ma_250": calc_ma(250),
                })
        
        info(f"Calculated {len(results)} quote indicators for {symbol}")
        return results
    
    async def fetch_kline_finance(
        self,
        symbol: str,
        start: str,
        end: str
    ) -> List[Dict[str, Any]]:
        """Fetch daily finance data
        
        Corresponding table: stock_klines_1d_finance
        Data sources: get_equity_structure, get_share_holder, get_income
        
        Returns:
            List[Dict] containing fields:
            - total_market_cap: Total market cap
            - float_market_cap: Float market cap
            - total_shares: Total shares
            - float_shares: Float shares
            - inst_holding_shares: Institutional holding shares
            - inst_holding_ratio: Institutional holding ratio
            - net_profit: Net profit
            - revenue: Revenue
            - eps: EPS
            - roe: ROE
        """
        print(f"[amazingdata_adapter fetch_kline_finance]Fetching {symbol} finance data...")
        self._check_login()
        
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            None,
            lambda: self._fetch_kline_finance_sync(symbol, start, end)
        )
    
    def _fetch_kline_finance_sync(
        self,
        symbol: str,
        start_date: str,
        end_date: str
    ) -> List[Dict[str, Any]]:
        """Sync method to fetch finance data"""
        codes = [symbol]
        start_int = self._format_date(start_date)
        end_int = self._format_date(end_date)
        
        results = []
        
        try:
            equity_dict = self._info_data.get_equity_structure(
                code_list=codes,
                local_path=self.config.local_path,
                is_local=self.config.use_local_cache
            )
            
            equity_data = {}
            if symbol in equity_dict:
                equity_df = equity_dict[symbol]
                for _, row in equity_df.iterrows():
                    ann_date = row.get('ANN_DATE')
                    if pd.notna(ann_date):
                        if isinstance(ann_date, (int, float)):
                            date_key = str(int(ann_date))
                        else:
                            date_key = str(ann_date).replace('-', '').replace('/', '')
                        equity_data[date_key] = {
                            'total_shares': float(row.get('TOT_A_SHARE', 0)) * 10000 if pd.notna(row.get('TOT_A_SHARE')) else 0,
                            'float_shares': float(row.get('FLOAT_A_SHARE', 0)) * 10000 if pd.notna(row.get('FLOAT_A_SHARE')) else 0,
                        }
        except Exception as e:
            print(f"[amazingdata_adapter]Failed to get equity structure: {e}")
            equity_data = {}
        
        kline_dict = self._market_data.query_kline(
            code_list=codes,
            begin_date=start_int,
            end_date=end_int,
            period=self._ad.constant.Period.day.value
        )
        
        if symbol not in kline_dict:
            return []
        
        df = kline_dict[symbol]
        
        for _, row in df.iterrows():
            kline_time = row.get('kline_time')
            if pd.isna(kline_time):
                continue
            
            if isinstance(kline_time, pd.Timestamp):
                trade_date = kline_time.strftime('%Y-%m-%d')
                trade_date_int = int(kline_time.strftime('%Y%m%d'))
            else:
                date_str = str(int(kline_time))
                if len(date_str) != 8:
                    continue
                dt = datetime.strptime(date_str, "%Y%m%d")
                trade_date = dt.strftime('%Y-%m-%d')
                trade_date_int = int(date_str)
            
            close = float(row.get('close', 0))
            
            total_shares = 0
            float_shares = 0
            for date_key in sorted(equity_data.keys(), reverse=True):
                if int(date_key) <= trade_date_int:
                    total_shares = equity_data[date_key]['total_shares']
                    float_shares = equity_data[date_key]['float_shares']
                    break
            
            total_market_cap = close * total_shares if total_shares > 0 and close > 0 else None
            float_market_cap = close * float_shares if float_shares > 0 and close > 0 else None
            
            results.append({
                "symbol": symbol,
                "trade_date": trade_date,
                "total_market_cap": round(total_market_cap, 2) if total_market_cap else None,
                "float_market_cap": round(float_market_cap, 2) if float_market_cap else None,
                "total_shares": int(total_shares) if total_shares > 0 else None,
                "float_shares": int(float_shares) if float_shares > 0 else None,
                "inst_holding_shares": None,
                "inst_holding_ratio": None,
                "top10_holders_ratio": None,
                "net_profit": None,
                "revenue": None,
                "eps": None,
                "roe": None,
                "trading_days": None,
            })
        
        info(f"Fetched {len(results)} finance records for {symbol}")
        return results
    
    async def fetch_stock_basic_info(
        self,
        codes: Optional[List[str]] = None
    ) -> List[Dict[str, Any]]:
        """Fetch stock basic info
        
        Corresponding table: stock_symbols
        Data source: get_stock_basic
        
        Returns:
            List[Dict] containing fields:
            - symbol_id: Symbol code
            - name: Name
            - exchange: Exchange
            - list_date: List date
            - list_board: List board
            - industry: Industry
            - status: Status
            - is_delisted: Is delisted
            - delist_date: Delist date
        """
        print(f"[amazingdata_adapter fetch_stock_basic_info]Fetching stock basic info...")
        self._check_login()
        
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            None,
            lambda: self._fetch_stock_basic_info_sync(codes)
        )
    
    def _fetch_stock_basic_info_sync(
        self,
        codes: Optional[List[str]] = None
    ) -> List[Dict[str, Any]]:
        """Sync method to fetch stock basic info"""
        try:
            all_codes = self._base_data.get_code_list(
                security_type=SecurityType.STOCK_A.value
            )
            
            if codes:
                all_codes = [c for c in all_codes if c in codes]
            
            info_df = self._base_data.get_code_info(
                security_type=SecurityType.STOCK_A.value
            )
            
            results = []
            for code in all_codes:
                if ".SH" in code:
                    exchange = "SH"
                elif ".SZ" in code:
                    exchange = "SZ"
                elif ".BJ" in code:
                    exchange = "BJ"
                else:
                    exchange = ""
                
                name = code
                if code in info_df.index and 'symbol' in info_df.columns:
                    name = info_df.loc[code, 'symbol']
                
                list_date = None
                try:
                    equity_dict = self._info_data.get_equity_structure(
                        code_list=[code],
                        local_path=self.config.local_path,
                        is_local=self.config.use_local_cache
                    )
                    if code in equity_dict and not equity_dict[code].empty:
                        first_record = equity_dict[code].iloc[0]
                        ann_date = first_record.get('ANN_DATE')
                        if pd.notna(ann_date):
                            if isinstance(ann_date, (int, float)):
                                list_date = datetime.strptime(str(int(ann_date)), "%Y%m%d")
                            else:
                                list_date = pd.to_datetime(ann_date)
                except:
                    pass
                
                results.append({
                    "symbol_id": code,
                    "name": name,
                    "exchange": exchange,
                    "list_date": list_date,
                    "list_board": None,
                    "industry": None,
                    "status": "active",
                    "is_delisted": False,
                    "delist_date": None,
                    "is_st": None,
                    "total_shares": None,
                    "float_shares": None,
                })
            
            info(f"Fetched {len(results)} stock basic info records")
            return results
            
        except Exception as e:
            error(f"Failed to fetch stock basic info: {e}")
            return []
'''

def main():
    # Read original file
    with open('app/adapters/amazingdata_adapter.py', 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Check if already added
    if 'fetch_kline_base' not in content:
        # Append new methods before the last line
        content = content.rstrip() + new_methods
        
        with open('app/adapters/amazingdata_adapter.py', 'w', encoding='utf-8') as f:
            f.write(content)
        print('New methods added successfully!')
    else:
        print('Methods already exist.')

if __name__ == "__main__":
    main()