stock_fundamentals/tools/stock_code_formatter.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
股票代码格式转换工具类

提供多种股票代码格式之间的转换功能
"""

import re
from typing import Optional


class StockCodeFormatter:
    """股票代码格式转换器"""
    
    def __init__(self):
        """初始化"""
        pass
    
    def _parse_stock_code(self, stock_code: str) -> tuple:
        """
        解析股票代码，提取代码和市场信息
        
        Args:
            stock_code: 输入的股票代码
            
        Returns:
            tuple: (代码数字, 市场标识) 如 ('688008', 'SH')
        """
        if not stock_code:
            return None, None
            
        stock_code = stock_code.strip().upper()
        
        # 处理 688008.SH 格式
        if '.' in stock_code:
            parts = stock_code.split('.')
            if len(parts) == 2:
                code, market = parts[0], parts[1]
                return code, market
        
        # 处理 SH688008 格式
        elif stock_code.startswith(('SZ', 'SH', 'BJ')):
            market = stock_code[:2]
            code = stock_code[2:]
            return code, market
        
        # 处理纯数字 688008 格式
        elif stock_code.isdigit():
            # 根据数字前缀自动判断市场
            if stock_code.startswith(('60', '68')):
                return stock_code, 'SH'
            elif stock_code.startswith(('00', '30', '20')):
                return stock_code, 'SZ'
            elif stock_code.startswith(('8', '43', '87')):
                return stock_code, 'BJ'
            else:
                return stock_code, None
        
        return None, None
    
    def to_dot_format(self, stock_code: str) -> Optional[str]:
        """
        转换为点分格式 (如: 688008.SH)
        
        Args:
            stock_code: 输入的股票代码，支持多种格式
            
        Returns:
            str: 点分格式的股票代码，如果无法转换则返回None
            
        Examples:
            >>> formatter = StockCodeFormatter()
            >>> formatter.to_dot_format('SH688008')
            '688008.SH'
            >>> formatter.to_dot_format('688008')
            '688008.SH'
            >>> formatter.to_dot_format('688008.SH')
            '688008.SH'
        """
        code, market = self._parse_stock_code(stock_code)
        
        if code and market:
            return f"{code}.{market}"
        
        return None
    
    def to_prefix_format(self, stock_code: str) -> Optional[str]:
        """
        转换为前缀格式 (如: SH688008)
        
        Args:
            stock_code: 输入的股票代码，支持多种格式
            
        Returns:
            str: 前缀格式的股票代码，如果无法转换则返回None
            
        Examples:
            >>> formatter = StockCodeFormatter()
            >>> formatter.to_prefix_format('688008.SH')
            'SH688008'
            >>> formatter.to_prefix_format('688008')
            'SH688008'
            >>> formatter.to_prefix_format('SH688008')
            'SH688008'
        """
        code, market = self._parse_stock_code(stock_code)
        
        if code and market:
            return f"{market}{code}"
        
        return None
    
    def to_number_only(self, stock_code: str) -> Optional[str]:
        """
        转换为纯数字格式 (如: 688008)
        
        Args:
            stock_code: 输入的股票代码，支持多种格式
            
        Returns:
            str: 纯数字格式的股票代码，如果无法转换则返回None
            
        Examples:
            >>> formatter = StockCodeFormatter()
            >>> formatter.to_number_only('688008.SH')
            '688008'
            >>> formatter.to_number_only('SH688008')
            '688008'
            >>> formatter.to_number_only('688008')
            '688008'
        """
        code, market = self._parse_stock_code(stock_code)
        
        if code:
            return code
        
        return None
    
    def get_market(self, stock_code: str) -> Optional[str]:
        """
        获取股票代码对应的市场标识
        
        Args:
            stock_code: 输入的股票代码，支持多种格式
            
        Returns:
            str: 市场标识 ('SH', 'SZ', 'BJ')，如果无法识别则返回None
        """
        code, market = self._parse_stock_code(stock_code)
        return market
    
    def is_valid_stock_code(self, stock_code: str) -> bool:
        """
        验证股票代码是否有效
        
        Args:
            stock_code: 输入的股票代码
            
        Returns:
            bool: 是否为有效的股票代码
        """
        code, market = self._parse_stock_code(stock_code)
        return code is not None and market is not None
    
    def batch_convert(self, stock_codes: list, target_format: str = 'dot') -> dict:
        """
        批量转换股票代码格式
        
        Args:
            stock_codes: 股票代码列表
            target_format: 目标格式 ('dot', 'prefix', 'number')
            
        Returns:
            dict: {原始代码: 转换后代码} 的映射字典
        """
        result = {}
        
        for stock_code in stock_codes:
            if target_format == 'dot':
                converted = self.to_dot_format(stock_code)
            elif target_format == 'prefix':
                converted = self.to_prefix_format(stock_code)
            elif target_format == 'number':
                converted = self.to_number_only(stock_code)
            else:
                converted = None
            
            result[stock_code] = converted
        
        return result


def main():
    """示例用法"""
    print("=== 股票代码格式转换工具示例 ===")
    
    # 创建格式转换器
    formatter = StockCodeFormatter()
    
    # 测试用例
    test_codes = [
        '688008.SH',   # 沪市科创板
        'SH688008',    # 前缀格式
        '688008',      # 纯数字
        '300661.SZ',   # 深市创业板
        'SZ300661',    # 前缀格式
        '300661',      # 纯数字
        '000858.SZ',   # 深市主板
        '600519.SH',   # 沪市主板
        '430123.BJ',   # 北交所
        'BJ430123',    # 北交所前缀格式
    ]
    
    print("\n原始代码 -> 点分格式 -> 前缀格式 -> 纯数字格式 -> 市场")
    print("-" * 70)
    
    for code in test_codes:
        dot_format = formatter.to_dot_format(code)
        prefix_format = formatter.to_prefix_format(code)
        number_format = formatter.to_number_only(code)
        market = formatter.get_market(code)
        is_valid = formatter.is_valid_stock_code(code)
        
        status = "✓" if is_valid else "✗"
        print(f"{code:12} -> {dot_format or 'None':12} -> {prefix_format or 'None':12} -> {number_format or 'None':8} -> {market or 'None':4} {status}")
    
    # 批量转换示例
    print(f"\n=== 批量转换示例 ===")
    batch_codes = ['SH688008', '300661', 'BJ430123']
    
    dot_results = formatter.batch_convert(batch_codes, 'dot')
    print("转换为点分格式:")
    for original, converted in dot_results.items():
        print(f"  {original} -> {converted}")
    
    prefix_results = formatter.batch_convert(batch_codes, 'prefix')
    print("转换为前缀格式:")
    for original, converted in prefix_results.items():
        print(f"  {original} -> {converted}")


if __name__ == "__main__":
    main()