Coverage for quantec/easydata/client.py: 75%

1import hashlib

2import logging

3import os

4from io import StringIO

5from pathlib import Path

6from typing import Optional, Union

8import pandas as pd

9import requests

10from dotenv import load_dotenv

12from .. import __version__

14load_dotenv()

16log = logging.getLogger(__name__)

19class Client:

20 """Client for Quantec API.

22 Parameters

23 ----------

24 api_key : Optional[str], optional

25 API key. Defaults to EASYDATA_API_KEY env variable.

26 api_url : Optional[str], optional

27 API base URL. Defaults to EASYDATA_API_URL env variable or https://www.easydata.co.za/api/v3/.

28 use_cache : bool, optional

29 Enable caching for grid data. Defaults to False.

30 cache_dir : str, optional

31 Directory for cached files. Defaults to 'cache'.

33 Raises

34 ------

35 ValueError

36 If api_key is empty.

38 """

40 def __init__(

41 self,

42 api_key: Optional[str] = None,

43 api_url: Optional[str] = None,

44 use_cache: bool = False,

45 cache_dir: str = "cache",

46 ) -> None:

47 api_key = api_key or os.getenv("EASYDATA_API_KEY")

48 api_url = api_url or os.getenv("EASYDATA_API_URL") or "https://www.easydata.co.za/api/v3"

49 if not api_key:

50 raise ValueError(

51 "API key must be provided via api_key parameter or EASYDATA_API_KEY environment variable"

52 )

54 self.__version__: str = __version__

55 self.api_key: str = api_key

56 self.api_url: str = api_url.rstrip("/")

57 self.use_cache: bool = use_cache

58 self.cache_dir: str = cache_dir

60 if use_cache:

61 self._setup_cache()

63 def get_data(

64 self,

65 time_series_codes: Optional[str] = None,

66 selection_pk: Optional[int] = None,

67 freq: str = "M",

68 start_year: str = "",

69 end_year: str = "",

70 analysis: bool = False,

71 resp_format: str = "csv",

72 is_tidy: bool = True,

73 ) -> Union[pd.DataFrame, dict]:

74 """

75 Fetch data from Quantec API.

77 Parameters

78 ----------

79 time_series_codes : Optional[str], optional

80 Comma-separated string of time series codes (e.g., "code1,code2").

81 selection_pk : Optional[int], optional

82 Selection primary key. Takes precedence over time_series_codes.

83 freq : str, optional

84 Data frequency ('M', 'Q', etc.). Defaults to 'M'.

85 start_year : str, optional

86 Start date ('YYYY-MM-DD'). Defaults to ''.

87 end_year : str, optional

88 End date ('YYYY-MM-DD'). Defaults to ''.

89 analysis : bool, optional

90 Include analysis parameter. Defaults to False.

91 resp_format : str, optional

92 Response format ('csv' or 'json'). Defaults to 'csv'.

93 is_tidy : bool, optional

94 Return tidy data. Defaults to True.

96 Returns

97 -------

98 Union[pd.DataFrame, dict]

99 DataFrame for CSV, dict for JSON.

100

101 Raises

102 ------

103 ValueError

104 If neither time_series_codes nor selection_pk is provided.

105 requests.HTTPError

106 If API request fails.

107 requests.ConnectionError

108 If network issue occurs.

109 ValueError

110 If response parsing fails.

111

112 """

113 if not time_series_codes and selection_pk is None:

114 raise ValueError(

115 "Either time_series_codes or selection_pk must be provided"

116 )

117

118 url: str = f"{self.api_url}/download/"

119

120 query_params: dict[str, Union[str, bool, int]] = {

121 "respFormat": resp_format,

122 "freqs": freq,

123 "startYear": start_year,

124 "endYear": end_year,

125 "isTidy": is_tidy,

126 "analysis": analysis,

127 }

128

129 if selection_pk is not None:

130 query_params["selectionPk"] = selection_pk

131 log_key = str(selection_pk)

132 else:

133 query_params["timeSeriesCodes"] = time_series_codes

134 log_key = time_series_codes

135

136 log.debug(f"[{log_key}] -- Querying with parameters: {query_params}")

137

138 try:

139 response = requests.get(

140 url, params={**query_params, "auth_token": self.api_key}

141 )

142 response.raise_for_status()

143 except requests.ConnectionError as e:

144 raise requests.ConnectionError(

145 "Network error: Unable to connect to API"

146 ) from e

147 except requests.HTTPError as e:

148 raise requests.HTTPError(f"API request failed: {response.text}") from e

149

150 if resp_format == "csv":

151 try:

152 out: pd.DataFrame = (

153 pd.read_csv(StringIO(response.text)).dropna().reset_index()

154 )

155 except pd.errors.ParserError as e:

156 raise ValueError("Failed to parse CSV response") from e

157 else:

158 try:

159 out: dict = response.json()

160 except ValueError as e:

161 raise ValueError("Failed to parse JSON response") from e

162

163 log.debug(f"[{log_key}] -- Found {len(out)} rows")

164 return out

165

166 def _setup_cache(self) -> None:

167 """Create cache directory if it doesn't exist."""

168 Path(self.cache_dir).mkdir(parents=True, exist_ok=True)

169

170 def _generate_cache_key(self, *args, debug: bool = False) -> str:

171 """Generate hash-based cache key from arguments."""

172 hash_input = "".join(str(arg) for arg in args)

173 if debug:

174 return f"debug_{hashlib.md5(hash_input.encode()).hexdigest()[:8]}"

175 return hashlib.md5(hash_input.encode()).hexdigest()

176

177 def _load_from_cache(

178 self, cache_key: str, resp_format: str

179 ) -> Optional[pd.DataFrame]:

180 """Load data from cache if it exists."""

181 if not self.use_cache:

182 return None

183

184 cache_path = Path(self.cache_dir) / f"{cache_key}.{resp_format}"

185 if not cache_path.exists():

186 return None

187

188 log.debug(f"Loading from cache: {cache_path}")

189 if resp_format == "parquet":

190 return pd.read_parquet(cache_path)

191 elif resp_format == "csv":

192 return pd.read_csv(cache_path)

193 return None

194

195 def _save_to_cache(

196 self, data: pd.DataFrame, cache_key: str, resp_format: str

197 ) -> None:

198 """Save data to cache."""

199 if not self.use_cache:

200 return

201

202 cache_path = Path(self.cache_dir) / f"{cache_key}.{resp_format}"

203 log.debug(f"Saving to cache: {cache_path}")

204 if resp_format == "parquet":

205 data.to_parquet(cache_path, index=False)

206 elif resp_format == "csv":

207 data.to_csv(cache_path, index=False)

208

209 def get_recipes(self) -> Union[pd.DataFrame, dict]:

210 """

211 Fetch available recipes from Quantec API.

212

213 Returns

214 -------

215 pd.DataFrame

216 DataFrame containing recipe information.

217

218 Raises

219 ------

220 requests.HTTPError

221 If API request fails.

222 requests.ConnectionError

223 If network issue occurs.

224 ValueError

225 If response parsing fails.

226

227 """

228 url: str = f"{self.api_url}/recipes/"

229

230 try:

231 response = requests.get(url, params={"auth_token": self.api_key})

232 response.raise_for_status()

233 except requests.ConnectionError as e:

234 raise requests.ConnectionError(

235 "Network error: Unable to connect to API"

236 ) from e

237 except requests.HTTPError as e:

238 raise requests.HTTPError(f"API request failed: {response.text}") from e

239

240 # Always return as DataFrame for recipes

241 try:

242 recipes_data = response.json()

243 out: pd.DataFrame = pd.DataFrame(recipes_data).dropna(axis=1, how="all")

244 except (ValueError, pd.errors.ParserError) as e:

245 raise ValueError("Failed to parse recipes response") from e

246

247 log.debug(

248 f"Found {len(out) if isinstance(out, pd.DataFrame) else len(out)} recipes"

249 )

250 return out

251

252 def get_selections(

253 self,

254 status: Optional[str] = None,

255 show: Optional[str] = None,

256 filter: Optional[str] = None,

257 ) -> Union[pd.DataFrame, dict]:

258 """

259 Fetch user's available selections from Quantec API.

260

261 Parameters

262 ----------

263 status : Optional[str], optional

264 Filter by selection status using combined flags:

265 U=Unsaved, P=Private, S=Shared, O=Open (e.g., "PSO").

266 show : Optional[str], optional

267 Show specific selection types ("shared" or "open").

268 filter : Optional[str], optional

269 Apply additional filters (e.g., "active").

270

271 Returns

272 -------

273 pd.DataFrame

274 Selection data with transformed fields: item, pk, title,

275 code_count, is_owner, owner, status, description, modified.

276

277 Raises

278 ------

279 requests.HTTPError

280 If API request fails.

281 requests.ConnectionError

282 If network issue occurs.

283 ValueError

284 If response parsing fails.

285

286 """

287 url: str = f"{self.api_url}/selections/"

288

289 query_params: dict[str, str] = {"auth_token": self.api_key, "format": "json"}

290

291 if status:

292 query_params["status"] = status

293 if show:

294 query_params["show"] = show

295 if filter:

296 query_params["filter"] = filter

297

298 log.debug(f"Querying selections with parameters: {query_params}")

299

300 try:

301 response = requests.get(url, params=query_params)

302 response.raise_for_status()

303 except requests.ConnectionError as e:

304 raise requests.ConnectionError(

305 "Network error: Unable to connect to API"

306 ) from e

307 except requests.HTTPError as e:

308 raise requests.HTTPError(f"API request failed: {response.text}") from e

309

310 try:

311 resp = response.json()

312 if not resp:

313 selections_data = []

314 else:

315 # Transform data following the original logic

316 selections_data = [

317 {

318 "item": i,

319 "pk": item["id"],

320 "title": item["title"],

321 "code_count": len(item.get("timeseriescodes", [])),

322 "is_owner": item["is_owner"],

323 "owner": item["owner"]["username"],

324 "status": item["status"],

325 "description": item.get("description", ""),

326 "modified": item["modified"],

327 }

328 for i, item in enumerate(resp, 1)

329 ]

330 except (ValueError, KeyError, TypeError) as e:

331 raise ValueError("Failed to parse selections response") from e

332

333 # Always return as DataFrame for selections

334 out: pd.DataFrame = pd.DataFrame(selections_data).dropna(axis=1, how="all")

335

336 log.debug(f"Found {len(selections_data)} selections")

337 return out

338

339 def get_grid_data(

340 self,

341 recipe_pk: int,

342 is_expanded: bool = True,

343 is_melted: bool = True,

344 resp_format: str = "dataframe",

345 selectdimensionnodes: dict = None,

346 ) -> Union[pd.DataFrame, str, bytes]:

347 """

348 Fetch grid/pivot table data using recipe primary key.

349

350 Parameters

351 ----------

352 recipe_pk : int

353 Recipe primary key identifier.

354 is_expanded : bool, optional

355 Return expanded data format. Defaults to True.

356 is_melted : bool, optional

357 Return melted data format. Defaults to True.

358 resp_format : str, optional

359 Response format ('dataframe', 'parquet', or 'csv'). Defaults to 'dataframe'.

360 selectdimensionnodes : dict, optional

361 Dimension filtering. Example: {"dimension": "d1", "codes": ["CODE1"]}.

362 Defaults to None.

363

364 Returns

365 -------

366 Union[pd.DataFrame, str, bytes]

367 DataFrame if resp_format='dataframe', CSV string if resp_format='csv',

368 or bytes if resp_format='parquet'.

369

370 Raises

371 ------

372 ValueError

373 If resp_format is invalid.

374 requests.HTTPError

375 If API request fails.

376 requests.ConnectionError

377 If network issue occurs.

378 ValueError

379 If response parsing fails.

380

381 """

382 if resp_format not in ["dataframe", "parquet", "csv"]:

383 raise ValueError("resp_format must be 'dataframe', 'parquet', or 'csv'")

384

385 # Determine API format (use parquet for dataframe requests for efficiency)

386 api_format = "parquet" if resp_format == "dataframe" else resp_format

387

388 # Check cache first

389 cache_key = self._generate_cache_key(

390 recipe_pk, is_expanded, is_melted, api_format, selectdimensionnodes

391 )

392

393 # Check if cached file exists and load raw data

394 if self.use_cache:

395 from pathlib import Path

396 cache_path = Path(self.cache_dir) / f"{cache_key}.{api_format}"

397 if cache_path.exists():

398 log.debug(f"[{recipe_pk}] -- Loading from cache: {cache_path}")

399

400 if resp_format == "csv":

401 # Return cached CSV string

402 return cache_path.read_text()

403 elif resp_format == "parquet":

404 # Return cached parquet bytes

405 return cache_path.read_bytes()

406 else: # resp_format == "dataframe"

407 # Parse cached file to DataFrame

408 if api_format == "parquet":

409 cached_df = pd.read_parquet(cache_path)

410 else: # api_format == "csv"

411 cached_df = pd.read_csv(cache_path)

412 return cached_df.dropna(axis=1, how="all")

413

414 url: str = f"{self.api_url}/download/recipes/{recipe_pk}/"

415

416 # Use POST if filtering, GET otherwise

417 if selectdimensionnodes:

418 # POST request for filtering

419 request_data = {

420 "respFormat": api_format,

421 "isExpanded": is_expanded,

422 "isMelted": is_melted,

423 "selectdimensionnodes": selectdimensionnodes,

424 }

425

426 headers = {

427 "Authorization": f"Token {self.api_key}",

428 "Content-Type": "application/json",

429 }

430

431 log.debug(f"[{recipe_pk}] -- POST with filters: {selectdimensionnodes}")

432

433 try:

434 response = requests.post(url, json=request_data, headers=headers)

435 response.raise_for_status()

436 except requests.ConnectionError as e:

437 raise requests.ConnectionError(

438 "Network error: Unable to connect to API"

439 ) from e

440 except requests.HTTPError as e:

441 raise requests.HTTPError(f"API request failed: {response.text}") from e

442 else:

443 # GET request (existing code)

444 query_params: dict[str, Union[str, bool, int]] = {

445 "respFormat": api_format,

446 "isExpanded": is_expanded,

447 "isMelted": is_melted,

448 "auth_token": self.api_key,

449 }

450

451 log.debug(f"[{recipe_pk}] -- Querying with parameters: {query_params}")

452

453 try:

454 response = requests.get(url, params=query_params)

455 response.raise_for_status()

456 except requests.ConnectionError as e:

457 raise requests.ConnectionError(

458 "Network error: Unable to connect to API"

459 ) from e

460 except requests.HTTPError as e:

461 raise requests.HTTPError(f"API request failed: {response.text}") from e

462

463 # Save raw response to cache first

464 if self.use_cache:

465 cache_path = Path(self.cache_dir) / f"{cache_key}.{api_format}"

466 log.debug(f"[{recipe_pk}] -- Saving to cache: {cache_path}")

467 if api_format == "parquet":

468 cache_path.write_bytes(response.content)

469 else: # api_format == "csv"

470 cache_path.write_text(response.text)

471

472 # Handle return format based on user's request

473 if resp_format == "csv":

474 # Return raw CSV string

475 log.debug(f"[{recipe_pk}] -- Returning raw CSV data")

476 return response.text

477 elif resp_format == "parquet":

478 # Return raw parquet bytes

479 log.debug(f"[{recipe_pk}] -- Returning raw parquet data")

480 return response.content

481 else: # resp_format == "dataframe"

482 # Parse into DataFrame and apply cleaning

483 if api_format == "parquet":

484 try:

485 from io import BytesIO

486 out: pd.DataFrame = pd.read_parquet(BytesIO(response.content))

487 except Exception as e:

488 raise ValueError("Failed to parse parquet response") from e

489 else: # api_format == "csv"

490 try:

491 out: pd.DataFrame = pd.read_csv(StringIO(response.text))

492 except pd.errors.ParserError as e:

493 raise ValueError("Failed to parse CSV response") from e

494

495 # Clean up data (only for DataFrame output)

496 out = out.dropna(axis=1, how="all")

497

498 log.debug(f"[{recipe_pk}] -- Found {len(out)} rows")

499 return out