Coverage for adrs/data/cache.py: 17%

1import logging 

2import polars as pl 

3 

4from pathlib import Path 

5from datetime import ( 

6 date, 

7 datetime, 

8 timedelta, 

9 timezone, 

10) 

11 

12from adrs.types import Topic 

13from adrs.data.datasource import Datasource 

14from adrs.data.utils import make_filename, make_filepath, parse_date_from_filename 

15 

16 

17logger = logging.getLogger(__name__) 

18 

19 

20def _generate_date_range(start: datetime, end: datetime) -> list[date]: 

21 dates, current = [], start.date() 

22 while current <= end.date(): 

23 dates.append(current) 

24 current += timedelta(days=1) 

25 return dates 

26 

27 

28class Cache: 

29 def __init__(self, data_path: Path, format: str, override_existing: bool): 

30 self.data_path: Path = data_path 

31 self.fmt: str = format 

32 self.override_existing: bool = override_existing 

33 

34 def init( 

35 self, topic: Topic, start_time: datetime, end_time: datetime 

36 ) -> tuple[datetime, datetime]: 

37 filepath = self.data_path / make_filepath(topic, self.fmt) 

38 parent_dir = filepath.parent 

39 if not parent_dir.exists(): 

40 logger.debug("[%s] creating directory at %s", topic, parent_dir) 

41 parent_dir.mkdir(parents=True, exist_ok=True) 

42 

43 # Floor start, ceil end to day boundaries 

44 start_time = start_time.replace(hour=0, minute=0, second=0, microsecond=0) 

45 end_midnight = end_time.replace(hour=0, minute=0, second=0, microsecond=0) 

46 end_time = ( 

47 end_midnight + timedelta(days=1) 

48 if end_time > end_midnight 

49 else end_midnight 

50 ) 

51 

52 if self.override_existing: 

53 logger.debug("[%s] override is set, skipping check for cache", topic) 

54 return (start_time, end_time) 

55 else: 

56 filename_base = make_filename(topic) 

57 downloaded_dates = [ 

58 d.date() 

59 for f in parent_dir.iterdir() 

60 if f.is_file() and filename_base in f.name 

61 if (d := parse_date_from_filename(f.name)) is not None 

62 ] 

63 

64 needed = _generate_date_range(start_time, end_time - timedelta(days=1)) 

65 filtered = [d for d in needed if d not in downloaded_dates] 

66 

67 if filtered: 

68 download_start_time = datetime.combine( 

69 min(filtered), datetime.min.time(), tzinfo=timezone.utc 

70 ) 

71 download_end_time = datetime.combine( 

72 max(filtered) + timedelta(days=1), 

73 datetime.min.time(), 

74 tzinfo=timezone.utc, 

75 ) 

76 else: 

77 download_start_time = end_time 

78 download_end_time = end_time 

79 return (download_start_time, download_end_time) 

80 

81 def _write_file(self, df: pl.DataFrame, path: Path, fmt: str) -> None: 

82 match fmt: 

83 case "parquet": 

84 df.write_parquet(path) 

85 case "json": 

86 df.write_json(path) 

87 case "csv": 

88 df.write_csv(path) 

89 case _: 

90 raise ValueError(f"unsupported format: {fmt}") 

91 

92 def _read_file(self, path: Path, fmt: str) -> pl.DataFrame: 

93 match fmt: 

94 case "parquet": 

95 return pl.read_parquet(path) 

96 case "json": 

97 return pl.read_json(path) 

98 case "csv": 

99 return pl.read_csv(path) 

100 case _: 

101 raise ValueError(f"unsupported format: {fmt}") 

102 

103 async def download( 

104 self, 

105 datasource: Datasource, 

106 topic: Topic, 

107 download_start_time: datetime, 

108 download_end_time: datetime, 

109 ) -> int: 

110 if download_start_time >= download_end_time: 

111 return 0 

112 logger.info( 

113 "[%s] downloading from %s to %s", 

114 topic, 

115 download_start_time, 

116 download_end_time, 

117 ) 

118 df = await datasource.query_paginated( 

119 topic=topic, 

120 start_time=download_start_time, 

121 end_time=download_end_time, 

122 flatten=True, 

123 ) 

124 if df.is_empty(): 

125 return 0 

126 

127 df = df.with_columns( 

128 pl.col("start_time").dt.strftime("%Y-%m-%d").alias("date"), 

129 *[ 

130 pl.col(col).cast(pl.Float64, strict=False) 

131 for col, dtype in df.schema.items() 

132 if dtype != pl.Float64 

133 and col not in ["start_time", "datetime", "date"] 

134 # Multi-record-per-timestamp topics (full option chains, GEX 

135 # snapshots) carry a nested List(Struct(...)) column -- not 

136 # castable to Float64 (structurally, not just a value-level 

137 # failure strict=False would relax). Leave nested columns 

138 # as-is; only cast genuinely scalar columns. 

139 and not isinstance(dtype, (pl.List, pl.Struct)) 

140 ], 

141 ) 

142 

143 for partition in df.partition_by("date", include_key=True): 

144 date_str = partition["date"][0] 

145 partition = partition.drop("date") 

146 

147 filepath = self.data_path / make_filepath(topic, self.fmt, suffix=date_str) 

148 filepath.parent.mkdir(parents=True, exist_ok=True) 

149 

150 logger.debug( 

151 "[%s] writing %d records to %s", topic, len(partition), filepath 

152 ) 

153 self._write_file(partition, filepath, self.fmt) 

154 

155 return len(df) 

156 

157 async def read( 

158 self, 

159 topic: Topic, 

160 start_time: datetime, 

161 end_time: datetime, 

162 ) -> pl.DataFrame: 

163 parent_dir = self.data_path / make_filepath(topic, self.fmt).parent 

164 if not parent_dir.exists(): 

165 raise FileNotFoundError(f"no files available at {parent_dir}") 

166 

167 filename_base = make_filename(topic) 

168 

169 entries = sorted( 

170 [f for f in parent_dir.iterdir() if f.is_file()], 

171 key=lambda f: parse_date_from_filename(f.name) 

172 or datetime.min.replace(tzinfo=timezone.utc), 

173 reverse=True, 

174 ) 

175 

176 frames: list[pl.DataFrame] = [] 

177 

178 for f in entries: 

179 date_dt = parse_date_from_filename(f.name) 

180 if ( 

181 date_dt is None 

182 or date_dt.date() < start_time.date() 

183 or date_dt >= end_time 

184 ): 

185 continue 

186 if filename_base not in f.name: 

187 continue 

188 

189 frames.append(self._read_file(f, self.fmt)) 

190 

191 if not frames: 

192 raise FileNotFoundError("no files read in the output dir") 

193 

194 return pl.concat(frames, how="diagonal_relaxed").rechunk().sort("start_time") 

195 

196 async def fetch( 

197 self, 

198 datasource: Datasource, 

199 topic: Topic, 

200 start_time: datetime, 

201 end_time: datetime, 

202 ): 

203 (download_start_time, download_end_time) = self.init( 

204 topic=topic, 

205 start_time=start_time, 

206 end_time=end_time, 

207 ) 

208 datapoints = await self.download( 

209 datasource=datasource, 

210 topic=topic, 

211 download_start_time=download_start_time, 

212 download_end_time=download_end_time, 

213 ) 

214 logger.info("[%s] downloaded %d datapoints", topic, datapoints) 

215 read_start = start_time.replace(hour=0, minute=0, second=0, microsecond=0) 

216 try: 

217 return await self.read( 

218 topic=topic, start_time=read_start, end_time=end_time 

219 ) 

220 except FileNotFoundError: 

221 return pl.DataFrame()