analyze_pattern.py

#!/usr/bin/env python3
import argparse
import asyncio
import logging
import os
import baostock as bs
import pandas as pd
import numpy as np
from tradingpatterns.tradingpatterns import (
    detect_head_shoulder, detect_multiple_tops_bottoms, 
    detect_triangle_pattern, detect_wedge, detect_channel, 
    detect_double_top_bottom, detect_trendline
)
from tradingpatterns.enum_types import TradingPattern
from tradingpatterns.pattern_utils import get_recent_bullish_patterns

logger = logging.getLogger("analyze_pattern")

def fetch_stock_data(stock_code, end_date, start_date="2022-01-01"):
    """
    Fetch stock data from BaoStock and convert to numeric DataFrame.
    """
    lg = bs.login()
    if lg.error_code != '0':
        logger.error(f"baostock login failed: {lg.error_msg}")
        return pd.DataFrame()

    rs = bs.query_history_k_data_plus(
        stock_code,
        "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",
        start_date=start_date, end_date=end_date,
        frequency="d", adjustflag="3"
    )
    
    if rs.error_code != '0':
        logger.error(f"baostock query failed: {rs.error_msg}")
        bs.logout()
        return pd.DataFrame()

    data_list = []
    while (rs.error_code == '0') & rs.next():
        data_list.append(rs.get_row_data())
    
    result = pd.DataFrame(data_list, columns=rs.fields)
    bs.logout()
    
    if not result.empty:
        # Map BaoStock columns to lowercase required by tradingpatterns
        # BaoStock returns open, high, low, close etc in lowercase already if specified in fields
        numeric_cols = ['open', 'high', 'low', 'close', 'volume', 'amount', 'turn', 'pctChg']
        for col in numeric_cols:
            if col in result.columns:
                result[col] = pd.to_numeric(result[col], errors='coerce')
    
    return result

async def main():
    # Configure logging
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    
    parser = argparse.ArgumentParser(description="Stock analysis command line tool using BaoStock and TradingPatternScanner")
    parser.add_argument("--stock_code", required=True, help="Stock code, e.g., sh.600000, sz.000001")
    parser.add_argument("--analysis_time", required=True, help="Analysis date, format: YYYY-MM-DD")
    
    args = parser.parse_args()
    
    logger.info(f"Analyzing stock code: {args.stock_code}, analysis date: {args.analysis_time}")
    
    # 判断本地是否存在，不存就下载
    if not os.path.exists(args.stock_code + ".csv"):
        df = fetch_stock_data(args.stock_code, args.analysis_time)
        df.to_csv(args.stock_code + ".csv", index=False)
    else:
        df = pd.read_csv(args.stock_code + ".csv")
    
    
    if df.empty:
        logger.error("No stock data retrieved.")
        return

    # Perform pattern detection
    logger.info(f"Retrieved {len(df)} rows of data. Columns: {df.columns.tolist()}")
    logger.debug(f"Data types:\n{df.dtypes}")
    
    # Define detection jobs
    jobs = [
        ("Head and Shoulder", detect_head_shoulder),
        ("Multiple Tops/Bottoms", detect_multiple_tops_bottoms),
        ("Triangle", detect_triangle_pattern),
        ("Wedge", detect_wedge),
        ("Channel", detect_channel),
        ("Double Top/Bottom", detect_double_top_bottom),
        ("Trendline", detect_trendline)
    ]
    
    for name, func in jobs:
        try:
            logger.info(f"Running {name} detection...")
            res_df = func(df)
            # Use concat to merge results back to df
            new_cols = res_df.columns.difference(df.columns)
            if not new_cols.empty:
                df = pd.concat([df, res_df[new_cols]], axis=1)
        except Exception as e:
            logger.error(f"Error during {name} detection: {e}", exc_info=True)
            # Continue to next job if one fails
    
    # Define categorical patterns (Enums) and numerical levels (Support/Resistance)
    pattern_cols = [
        'head_shoulder_pattern', 'multiple_top_bottom_pattern', 
        'triangle_pattern', 'wedge_pattern', 'channel_pattern', 
        'double_pattern'
    ]
    #level_cols = ['support', 'resistance']
    
    found_any = False
    n = 5
    recent_df = df.tail(n)
    
    # Report categorical patterns
    for col in pattern_cols:
        if col in df.columns:
            matches = recent_df[recent_df[col].notna()]
            if not matches.empty:
                for idx, row in matches.iterrows():
                    logger.info(f"Pattern detected on {row['date']}: [{col}] -> {row[col]}")
                found_any = True
                
    # Report numerical levels
    # for col in level_cols:
    #     if col in df.columns:
    #         matches = recent_df[recent_df[col].notna()]
    #         if not matches.empty:
    #             for idx, row in matches.iterrows():
    #                 logger.info(f"Technical Level on {row['date']}: {col.capitalize()} -> {row[col]:.2f}")
    #             found_any = True
    
    # Report recent bullish patterns
    bullish_patterns = get_recent_bullish_patterns(df, n)
    if bullish_patterns:
        logger.info(f"Summary of bullish patterns in the last {n} days:")
        for bp in bullish_patterns:
            logger.info(f"  - {bp['date']}: {bp['pattern']}")
    else:
        logger.info(f"No bullish patterns detected in the last {n} days.")

    if not found_any:
        logger.info(f"No significant patterns detected in the last {n} days.")

if __name__ == "__main__":
    asyncio.run(main())