Coverage for src / monte_neo / data / downloader.py: 99%

1"""Binance data downloader module. 

2 

3Downloads OHLCV data from Binance API and saves in Parquet format. 

4""" 

5 

6from __future__ import annotations 

7 

8import os 

9import time 

10from collections.abc import Callable 

11from datetime import datetime, timedelta 

12 

13import pandas as pd 

14from binance.spot import Spot 

15 

16from monte_neo.utils.logger import get_logger 

17 

18logger = get_logger(__name__) 

19 

20 

21class BinanceDownloader: 

22 """Download OHLCV data from Binance.""" 

23 

24 TIMEFRAME_MAP = { 

25 "1m": "1m", 

26 "5m": "5m", 

27 "15m": "15m", 

28 "30m": "30m", 

29 "1h": "1h", 

30 "4h": "4h", 

31 "1d": "1d", 

32 "1w": "1w", 

33 } 

34 

35 COLUMNS = [ 

36 "timestamp", 

37 "open", 

38 "high", 

39 "low", 

40 "close", 

41 "volume", 

42 "close_time", 

43 "quote_volume", 

44 "trades", 

45 "taker_buy_base", 

46 "taker_buy_quote", 

47 "ignore", 

48 ] 

49 

50 def __init__( 

51 self, 

52 api_key: str | None = None, 

53 api_secret: str | None = None, 

54 ) -> None: 

55 """Initialize Binance client. 

56 

57 Args: 

58 api_key: Optional Binance API key. 

59 api_secret: Optional Binance API secret. 

60 """ 

61 self.api_key = api_key or os.getenv("BINANCE_API_KEY", "") 

62 self.api_secret = api_secret or os.getenv("BINANCE_API_SECRET", "") 

63 self.client = Spot(api_key=self.api_key, api_secret=self.api_secret) 

64 logger.info("Binance client initialized") 

65 

66 def get_available_symbols(self) -> list[str]: 

67 """Get list of available trading symbols. 

68 

69 Returns: 

70 List of symbol strings (e.g., ['BTCUSDT', 'ETHUSDT']). 

71 """ 

72 info = self.client.exchange_info() 

73 return sorted( 

74 [s["symbol"] for s in info["symbols"] if s["status"] == "TRADING"] 

75 ) 

76 

77 def _fetch_klines( 

78 self, 

79 symbol: str, 

80 interval: str, 

81 start_ms: int, 

82 end_ms: int, 

83 ) -> list[list[str | int | float]]: 

84 klines: list[list[str | int | float]] = [] 

85 current_start = start_ms 

86 while current_start < end_ms: 

87 try: 

88 batch = self.client.klines( 

89 symbol=symbol, 

90 interval=interval, 

91 startTime=current_start, 

92 endTime=end_ms, 

93 limit=1000, 

94 ) 

95 if not batch: 

96 break 

97 klines.extend(batch) 

98 last_close = int(batch[-1][6]) 

99 next_start = last_close + 1 

100 

101 # Compliance with Binance rate limits (small delay between batches) 

102 if len(batch) >= 1000: 

103 time.sleep(0.1) # 100ms delay between 1000-candle batches 

104 

105 if next_start <= current_start: 

106 break 

107 current_start = next_start 

108 except Exception as e: 

109 if "429" in str(e) or "rate limit" in str(e).lower(): 

110 logger.warning("Rate limit hit, sleeping for 10 seconds...") 

111 time.sleep(10) 

112 continue 

113 raise e 

114 return klines 

115 

116 def download( 

117 self, 

118 symbol: str, 

119 timeframe: str, 

120 start_date: datetime | str, 

121 end_date: datetime | str | None = None, 

122 progress_callback: Callable[[int, int, str], None] | None = None, 

123 ) -> pd.DataFrame: 

124 """Download OHLCV data from Binance. 

125 

126 Args: 

127 symbol: Trading pair symbol (e.g., 'BTCUSDT'). 

128 timeframe: Candle interval (e.g., '1h', '4h', '1d'). 

129 start_date: Start date for data. 

130 end_date: End date for data (default: now). 

131 

132 Returns: 

133 DataFrame with OHLCV data. 

134 """ 

135 if isinstance(start_date, str): 

136 start_date = datetime.fromisoformat(start_date) 

137 if end_date is None: 

138 end_date = datetime.now() 

139 elif isinstance(end_date, str): 

140 end_date = datetime.fromisoformat(end_date) 

141 

142 interval = self.TIMEFRAME_MAP.get(timeframe) 

143 if interval is None: 

144 raise ValueError(f"Invalid timeframe: {timeframe}") 

145 

146 if progress_callback: 

147 progress_callback(0, 100, "Initializing...") 

148 

149 chunks = [] 

150 current_start = start_date 

151 total_days = (end_date - start_date).days or 1 

152 

153 while current_start < end_date: 

154 current_end = min(current_start + timedelta(days=30), end_date) 

155 

156 logger.debug(f"Fetching chunk: {current_start} to {current_end}") 

157 if progress_callback: 

158 pct = int(((current_start - start_date).days / total_days) * 100) 

159 progress_callback( 

160 pct, 100, f"Fetching {current_start.strftime('%Y-%m')}" 

161 ) 

162 

163 start_ms = int(current_start.timestamp() * 1000) 

164 end_ms = int(current_end.timestamp() * 1000) 

165 klines = self._fetch_klines( 

166 symbol=symbol, 

167 interval=interval, 

168 start_ms=start_ms, 

169 end_ms=end_ms, 

170 ) 

171 chunks.extend(klines) 

172 current_start = current_end + timedelta(milliseconds=1) 

173 

174 if progress_callback: 

175 progress_callback(100, 100, "Processing...") 

176 

177 df = pd.DataFrame(chunks, columns=self.COLUMNS) 

178 df = self._process_dataframe(df) 

179 

180 if progress_callback: 

181 progress_callback(100, 100, "Done") 

182 

183 logger.info(f"Downloaded {len(df)} candles") 

184 return df 

185 

186 def _process_dataframe(self, df: pd.DataFrame) -> pd.DataFrame: 

187 """Process raw Binance data. 

188 

189 Args: 

190 df: Raw DataFrame from Binance API. 

191 

192 Returns: 

193 Processed DataFrame with proper types. 

194 """ 

195 # Convert timestamp to datetime 

196 df["timestamp"] = pd.to_datetime(df["timestamp"], unit="ms") 

197 df.set_index("timestamp", inplace=True) 

198 

199 # Convert price and volume columns to float 

200 price_cols = ["open", "high", "low", "close", "volume", "quote_volume"] 

201 for col in price_cols: 

202 df[col] = df[col].astype(float) 

203 

204 # Keep only essential columns 

205 df = df[["open", "high", "low", "close", "volume"]] 

206 

207 return df 

208 

209 def download_multiple( 

210 self, 

211 symbols: list[str], 

212 timeframe: str, 

213 start_date: datetime | str, 

214 end_date: datetime | str | None = None, 

215 ) -> dict[str, pd.DataFrame]: 

216 """Download data for multiple symbols. 

217 

218 Args: 

219 symbols: List of trading pair symbols. 

220 timeframe: Candle interval. 

221 start_date: Start date for data. 

222 end_date: End date for data. 

223 

224 Returns: 

225 Dictionary mapping symbols to DataFrames. 

226 """ 

227 result = {} 

228 for symbol in symbols: 

229 try: 

230 result[symbol] = self.download(symbol, timeframe, start_date, end_date) 

231 except Exception as e: 

232 logger.error(f"Failed to download {symbol}: {e}") 

233 return result 

234 

235 

236def download_sample_data( 

237 symbol: str = "BTCUSDT", 

238 timeframe: str = "1h", 

239 days: int = 365, 

240) -> pd.DataFrame: 

241 """Convenience function to download sample data. 

242 

243 Args: 

244 symbol: Trading pair symbol. 

245 timeframe: Candle interval. 

246 days: Number of days of historical data. 

247 

248 Returns: 

249 DataFrame with OHLCV data. 

250 """ 

251 downloader = BinanceDownloader() 

252 end_date = datetime.now() 

253 start_date = end_date - timedelta(days=days) 

254 return downloader.download(symbol, timeframe, start_date, end_date)