Coverage for brokers / korea_investment / korea_invest_trid_provider.py: 95%

1from __future__ import annotations 

2from typing import Dict, Optional, Union 

3from enum import Enum 

4 

5from config.config_loader import load_configs, load_config, TR_IDS_CONFIG_PATH 

6from brokers.korea_investment.korea_invest_env import KoreaInvestApiEnv 

7from brokers.korea_investment.korea_invest_trid_keys import TrIdLeaf, TrId 

8 

9_DOMAINS = ("quotations", "account", "trading") 

10 

11class KoreaInvestTrIdProvider: 

12 """ 

13 - tr_ids_config.yaml에서 TR-ID 테이블을 읽어와 키 기반으로 안전하게 조회 

14 - 모드(모의/실전)에 따라 자동으로 paper/real leaf 선택(논리 키 TrId 사용 시) 

15 - 문자열/Enum 모두 지원 

16 """ 

17 

18 def __init__(self, env: KoreaInvestApiEnv, tr_ids: Dict) -> None: 

19 self._env = env 

20 if not isinstance(tr_ids, dict) or not tr_ids: 

21 raise ValueError("tr_ids 설정이 비어 있거나 올바르지 않습니다.") 

22 self._tr_ids = tr_ids # shape: {'quotations': {...}, 'account': {...}, 'trading': {...}} 

23 

24 # ---------- 팩토리 ---------- 

25 @classmethod 

26 def from_config_loader(cls, env: KoreaInvestApiEnv) -> "KoreaInvestTrIdProvider": 

27 merged = load_configs() 

28 tr_ids = merged.get("tr_ids") 

29 if not tr_ids: 

30 # 병합에 없다면 원본 파일 직접 로드 

31 tr_ids = load_config(TR_IDS_CONFIG_PATH).get("tr_ids") 

32 return cls(env, tr_ids) 

33 

34 # ---------- 내부 유틸 ---------- 

35 def _get_leaf_value(self, leaf_key: str) -> str: 

36 """ 

37 leaf_key가 어느 도메인(account/quotations/trading)에 있든 값을 찾아 반환. 

38 """ 

39 for domain in _DOMAINS: 

40 dom = self._tr_ids.get(domain, {}) 

41 if leaf_key in dom: 

42 return dom[leaf_key] 

43 raise KeyError(f"tr_ids에 '{leaf_key}'를 찾을 수 없습니다.") 

44 

45 # ---------- 공개 API ---------- 

46 def get_by_leaf(self, key: Union[str, TrIdLeaf]) -> str: 

47 """leaf 키로 직접 조회 (모드 무시)""" 

48 leaf = key.value if isinstance(key, Enum) else str(key) 

49 return self._get_leaf_value(leaf) 

50 

51 def get(self, key: Union[str, TrId, TrIdLeaf], **kwargs) -> str: 

52 """ 

53 - TrId(논리키): 모드에 따라 paper/real leaf 자동 선택 

54 - TrIdLeaf/str(leaf키): 그대로 반환 

55 """ 

56 # 논리 키 처리 

57 if isinstance(key, TrId): 57 ↛ 71line 57 didn't jump to line 71 because the condition on line 57 was always true

58 is_paper = bool(self._env.is_paper_trading) 

59 if key is TrId.INQUIRE_BALANCE: 

60 leaf = TrIdLeaf.INQUIRE_BALANCE_PAPER if is_paper else TrIdLeaf.INQUIRE_BALANCE_REAL 

61 return self.get_by_leaf(leaf) 

62 if key is TrId.ORDER_CASH_BUY: 

63 leaf = TrIdLeaf.ORDER_CASH_BUY_PAPER if is_paper else TrIdLeaf.ORDER_CASH_BUY_REAL 

64 return self.get_by_leaf(leaf) 

65 if key is TrId.ORDER_CASH_SELL: 65 ↛ 68line 65 didn't jump to line 68 because the condition on line 65 was always true

66 leaf = TrIdLeaf.ORDER_CASH_SELL_PAPER if is_paper else TrIdLeaf.ORDER_CASH_SELL_REAL 

67 return self.get_by_leaf(leaf) 

68 raise KeyError(f"지원하지 않는 논리 키: {key}") 

69 

70 # leaf 키 그대로 

71 return self.get_by_leaf(key) 

72 

73 # 편의 함수(읽기 쉬운 API) 

74 def quotations(self, leaf: TrIdLeaf) -> str: 

75 return self.get_by_leaf(leaf) 

76 

77 def account_inquire_balance(self) -> str: 

78 return self.get(TrId.INQUIRE_BALANCE) 

79 

80 def trading_order_cash(self, is_buy: bool) -> str: 

81 return self.get(TrId.ORDER_CASH_BUY if is_buy else TrId.ORDER_CASH_SELL) 

82 

83 def daily_itemchartprice(self) -> str: 

84 leaf = TrIdLeaf.DAILY_ITEMCHARTPRICE 

85 return self.get_by_leaf(leaf) 

86 

87 def time_itemchartprice(self) -> str: 

88 leaf = TrIdLeaf.TIME_ITEMCHARTPRICE 

89 return self.get_by_leaf(leaf) 

90 

91 def time_daily_itemchartprice(self) -> str: 

92 leaf = TrIdLeaf.TIME_DAILY_ITEMCHARTPRICE 

93 return self.get_by_leaf(leaf)