chart_analysis.py

import pymongo
from pymongo import MongoClient
from datetime import datetime
import os



def get_stock_data(stock_code: str, period: str):
    """
    MongoDB에서 특정 종목의 시세 데이터를 조회합니다.

    Args:
        stock_code (str): 조회할 종목의 코드 (예: "005930")
        period (str): 데이터 주기. "D"(일봉), "W"(주봉), "M"(월봉) 중 하나.

    Returns:
        dict: 'data' 키에는 캔들 데이터 리스트, 'count' 키에는 데이터의 개수를 담은 딕셔너리.
              데이터가 없거나 period가 잘못된 경우, 'data'는 빈 리스트, 'count'는 0이 됩니다.
    """
    # 1. 입력값 유효성 검사
    if period not in ["D", "W", "M"]:
        print(f"오류: period는 'D', 'W', 'M' 중 하나여야 합니다. (입력값: {period})")
        return {"data": [], "count": 0}

    MAX_COUNT = 120  # 가져올 데이터의 최대 개수

    # 2. MongoDB 쿼리 실행
    # 최신 데이터부터 가져오기 위해 날짜(date) 기준 내림차순(DESCENDING)으로 정렬
    query = {"stock_code": stock_code, "period": period}

    MONGO_URI = os.getenv("ATLAS_MONGO_URI")
    DB_NAME = "BarbellAI"
    COLLECTION_NAME = "stock"

    # MongoDB에 한 번만 연결합니다.
    client = pymongo.MongoClient(MONGO_URI)
    db = client[DB_NAME]
    collection = db[COLLECTION_NAME]

    try:
        cursor = collection.find(query).sort("date", pymongo.DESCENDING).limit(MAX_COUNT)
        
        # MongoDB의 문서를 리스트 형태로 변환
        raw_data_list = list(cursor)

    except Exception as e:
        print(f"DB 조회 중 오류 발생: {e}")
        return {"data": [], "count": 0}


    # 3. 데이터가 없는 경우 처리
    if not raw_data_list:
        return {"data": [], "count": 0}

    # 4. 차트 표기를 위해 시간 순서대로 재정렬 (오래된 날짜 -> 최신 날짜)
    raw_data_list.reverse()

    # 5. 최종 결과 포맷팅
    candles = []
    for doc in raw_data_list:
        candles.append({
            "date": doc["date"].strftime("%Y-%m-%d"),
            "open": doc.get("open"),
            "high": doc.get("high"),
            "low": doc.get("low"),
            "close": doc.get("close"),
            "volume": doc.get("volume")
        })

    # 6. 최종 결과 반환
    return {
        "data": candles,
        "count": len(candles)  # 실제 가져온 데이터의 개수
    }

print(get_stock_data("005930", "W"))