Source code for pyqqq.utils.futures_code

"""지수선물 근월물(최근월물) 단축코드 생성 유틸리티"""

import datetime as dtm

# 상품별 단축코드 구성 요소. xy는 프리픽스의 뒤 2자리 (구코드 1XY / 신코드 AXY 공통),
# quarterly는 분기월물(3/6/9/12월)만 상장되는지 여부
PRODUCTS = {
    "코스피200": {"xy": "01", "quarterly": True},
    "미니코스피200": {"xy": "05", "quarterly": False},
    "코스닥150": {"xy": "06", "quarterly": True},
}

# 구코드 연도문자. 코드 체계는 조회 날짜가 아니라 "만기 연도"로 결정된다:
# 만기 2025년 이하 월물 = 구코드만, 2026년 이상 월물 = 신코드만 존재
OLD_YEAR_CHARS = {2024: "V", 2025: "W"}


def second_thursday(year: int, month: int) -> dtm.date:
    """만기일: 결제월의 두 번째 목요일"""
    first = dtm.date(year, month, 1)
    return first + dtm.timedelta(days=(3 - first.weekday()) % 7 + 7)


def front_month(date: dtm.date, quarterly: bool) -> tuple:
    """조회 일자의 최근월물 (year, month). 만기일 당일까지는 해당 월물을 유지한다."""
    y, m = date.year, date.month
    while True:
        if (not quarterly or m in (3, 6, 9, 12)) and date <= second_thursday(y, m):
            return y, m
        m += 1
        if m > 12:
            m, y = 1, y + 1




def get_front_month_futures_code(product: str, date: dtm.date) -> str:
    """
    지수선물 상품의 특정 일자 기준 근월물(최근월물) 단축코드를 생성합니다.

    만기일(결제월 두 번째 목요일) 당일까지는 해당 월물 코드를 유지하고, 다음 날부터 다음 월물로 넘어갑니다.
    코드 체계는 월물의 만기 연도로 결정됩니다. (만기 2025년 이하 = 구코드 1XY..., 2026년 이상 = 신코드 AXY...)

    Args:
        product (str): 상품 이름. 코스피200, 미니코스피200, 코스닥150 중 하나
        date (datetime.date): 기준 일자

    Returns:
        str: 근월물 단축코드 (예: "A01609", "101W09")

    Raises:
        ValueError: 지원하지 않는 상품 이름인 경우

    Examples:
        >>> get_front_month_futures_code("코스피200", datetime.date(2026, 8, 5))
        'A01609'
        >>> get_front_month_futures_code("미니코스피200", datetime.date(2026, 8, 5))
        'A05608'
        >>> get_front_month_futures_code("코스피200", datetime.date(2025, 7, 21))
        '101W09'
    """
    assert type(date) is dtm.date, "date must be a datetime.date object"

    if product not in PRODUCTS:
        raise ValueError(f"지원하지 않는 상품입니다: {product} (코스피200, 미니코스피200, 코스닥150 중 하나여야 합니다)")

    entry = PRODUCTS[product]
    year, month = front_month(date, entry["quarterly"])

    if year in OLD_YEAR_CHARS:
        return f"1{entry['xy']}{OLD_YEAR_CHARS[year]}{month:02d}"
    return f"A{entry['xy']}{year % 10}{month:02d}"