예제 소스

build_stock_list.py

# -*- coding: utf-8 -*-
"""
FinanceDataReader 를 이용해 KOSPI/KOSDAQ 전종목 정보를 조회하여
'투자 유의 종목'(관리종목 / 투자유의 / 투자환기 / 정리매매 / 단기과열 등),
우선주(종목명 …우 / 2우B / 4우(전환) 등), ETF/ETN/ELW/리츠 등 비주식을 제외한
종목 리스트를 CSV 파일로 저장합니다.

사용법
  $ python build_stock_list.py

출력
  stock_list.csv
    - 종목코드, 종목명, 시장, 소속부

의존성
  pip install finance-datareader pandas
"""

from __future__ import annotations

import re
import sys
from pathlib import Path

import pandas as pd

try:
    import FinanceDataReader as fdr
except ImportError as exc:  # pragma: no cover
    print(
        "FinanceDataReader 가 설치되어 있지 않습니다. "
        "`pip install finance-datareader` 로 먼저 설치해 주세요.",
        file=sys.stderr,
    )
    raise SystemExit(1) from exc


# 소속부(Dept) 중 '투자 유의' 로 간주하여 제외할 키워드.
# FDR 의 Dept 컬럼은 여러 사유가 콤마로 같이 들어오는 경우가 있어
# 부분 일치(키워드 포함) 로 판정합니다.
# 종목명 끝 패턴: 삼성전자우, 현대차2우B, CJ4우(전환), SK우 등
PREFERRED_NAME_RE = re.compile(
    r"(?:우선주|[1-5]?우[BC]?|\d+우)(?:\([^)]*\))?$"
)

EXCLUDE_DEPT_KEYWORDS: tuple[str, ...] = (
    "관리종목",
    "투자유의",
    "투자환기",
    "투자주의",
    "투자경고",
    "투자위험",
    "정리매매",
    "단기과열",
    "거래정지",
    "불성실공시",
)


def fetch_korean_stocks() -> pd.DataFrame:
    """KOSPI/KOSDAQ 전종목 DataFrame 을 반환합니다.

    반환 컬럼(주요): Code, Name, Market, Dept
    """
    frames: list[pd.DataFrame] = []
    for market in ("KOSPI", "KOSDAQ"):
        df = fdr.StockListing(market)
        if df is None or df.empty:
            continue
        # Market 컬럼이 없으면 직접 채워줌
        if "Market" not in df.columns:
            df = df.assign(Market=market)
        frames.append(df)

    if not frames:
        return pd.DataFrame(columns=["Code", "Name", "Market", "Dept"])

    merged = pd.concat(frames, ignore_index=True)

    # 컬럼 정규화: 일부 버전에서는 Dept 가 없을 수 있음
    if "Dept" not in merged.columns:
        merged["Dept"] = ""

    keep_cols = [c for c in ("Code", "Name", "Market", "Dept") if c in merged.columns]
    return merged[keep_cols].copy()


def is_preferred_stock(name: object, code: object = None) -> bool:
    """우선주 여부 (종목명 기준). code 는 예비용으로만 받으며 현재는 name 으로 판정."""
    _ = code
    if not isinstance(name, str):
        return False
    n = name.strip()
    if not n:
        return False
    if "우선주" in n:
        return True
    return bool(PREFERRED_NAME_RE.search(n))


def _is_excluded_dept(dept: object) -> bool:
    if not isinstance(dept, str) or not dept.strip():
        return False
    return any(kw in dept for kw in EXCLUDE_DEPT_KEYWORDS)


def _is_stock_code(code: object) -> bool:
    """6자리 숫자 종목코드만 일반 주권으로 간주 (ETF/ETN/ELW 등 일부 제외)."""
    if not isinstance(code, str):
        return False
    return len(code) == 6 and code.isdigit()


def filter_korean_stocks(df: pd.DataFrame) -> tuple[pd.DataFrame, dict[str, int]]:
    stats = {
        "total": int(len(df)),
        "non_stock_skipped": 0,
        "section_excluded": 0,
        "preferred_excluded": 0,
    }

    if df.empty:
        return df.copy(), stats

    code_mask = df["Code"].map(_is_stock_code)
    stats["non_stock_skipped"] = int((~code_mask).sum())

    dept_mask = df["Dept"].map(_is_excluded_dept) if "Dept" in df.columns else False
    stats["section_excluded"] = int(pd.Series(dept_mask).sum()) if "Dept" in df.columns else 0

    if "Name" in df.columns:
        pref_mask = df.apply(
            lambda r: is_preferred_stock(r.get("Name"), r.get("Code")),
            axis=1,
        )
    else:
        pref_mask = False
    stats["preferred_excluded"] = int(pd.Series(pref_mask).sum()) if "Name" in df.columns else 0

    keep_mask = (
        code_mask
        & (~pd.Series(dept_mask, index=df.index) if "Dept" in df.columns else True)
        & (~pd.Series(pref_mask, index=df.index) if "Name" in df.columns else True)
    )
    out = df.loc[keep_mask].copy()

    out = out.rename(
        columns={
            "Code": "종목코드",
            "Name": "종목명",
            "Market": "시장",
            "Dept": "소속부",
        }
    )
    out["소속부"] = out.get("소속부", "").fillna("")
    out = out.drop_duplicates(subset=["종목코드"]).sort_values("종목코드").reset_index(drop=True)
    return out, stats


def main() -> int:
    out_path = Path(__file__).resolve().parent / "stock_list.csv"

    print("[1/3] 국내 주식 목록 조회 중 (KOSPI/KOSDAQ)...")
    try:
        all_stocks_df = fetch_korean_stocks()
    except Exception as exc:
        print(f"  ✗ 전종목 데이터 조회 실패: {exc}", file=sys.stderr)
        return 1

    if all_stocks_df.empty:
        print("  ✗ 데이터 조회 결과가 비어 있습니다.", file=sys.stderr)
        return 1

    print(f"  · 조회된 총 후보 종목 수: {len(all_stocks_df)}")

    print("[2/3] 필터링 중...")
    filtered_df, stats = filter_korean_stocks(all_stocks_df)

    print(f"  · 총 종목 수      : {stats['total']}")
    print(f"  · 비주식 제외     : {stats['non_stock_skipped']}")
    print(f"  · 유의/관리 제외  : {stats['section_excluded']}")
    print(f"  · 우선주 제외     : {stats['preferred_excluded']}")
    print(f"  · 최종 유효 종목수: {len(filtered_df)}")

    print(f"[3/3] CSV 저장: {out_path}")
    filtered_df.to_csv(out_path, encoding="utf-8-sig", index=False)
    print(f"  ✓ 저장 완료 ({len(filtered_df)} 종목)")
    return 0


if __name__ == "__main__":
    sys.exit(main())