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

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 apikey : Optional[str], optional

25 API key. Defaults to EASYDATA_API_KEY env variable.

26 respformat : str, optional

27 Response format ('csv' for time series, 'csv'/'parquet' for grid data). Defaults to 'csv'.

28 is_tidy : bool, optional

29 Return tidy data. Defaults to True.

30 api_url : Optional[str], optional

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

32 use_cache : bool, optional

33 Enable caching for grid data. Defaults to False.

34 cache_dir : str, optional

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

37 Raises

38 ------

39 ValueError

40 If apikey is empty or respformat is invalid.

42 """

44 def __init__(

45 self,

46 apikey: Optional[str] = None,

47 respformat: str = "csv",

48 is_tidy: bool = True,

49 api_url: Optional[str] = None,

50 use_cache: bool = False,

51 cache_dir: str = "cache",

52 ) -> None:

53 apikey = apikey or os.getenv("EASYDATA_API_KEY")

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

55 if not apikey:

56 raise ValueError(

57 "API key must be provided via apikey parameter or EASYDATA_API_KEY environment variable"

58 )

59 if respformat not in ["csv", "json", "parquet"]:

60 raise ValueError("respformat must be 'csv', 'json', or 'parquet'")

62 self.__version__: str = __version__

63 self.apikey: str = apikey

64 self.respformat: str = respformat

65 self.is_tidy: bool = is_tidy

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

67 self.use_cache: bool = use_cache

68 self.cache_dir: str = cache_dir

70 if use_cache:

71 self._setup_cache()

73 def get_data(

74 self,

75 time_series_codes: Optional[str] = None,

76 selection_pk: Optional[int] = None,

77 freq: str = "M",

78 start_year: str = "",

79 end_year: str = "",

80 analysis: bool = False,

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

82 """

83 Fetch data from Quantec API.

85 Parameters

86 ----------

87 time_series_codes : Optional[str], optional

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

89 selection_pk : Optional[int], optional

90 Selection primary key. Takes precedence over time_series_codes.

91 freq : str, optional

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

93 start_year : str, optional

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

95 end_year : str, optional

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

97 analysis : bool, optional

98 Include analysis parameter. Defaults to False.

100 Returns

101 -------

102 Union[pd.DataFrame, dict]

103 DataFrame for CSV, dict for JSON.

104

105 Raises

106 ------

107 ValueError

108 If neither time_series_codes nor selection_pk is provided.

109 requests.HTTPError

110 If API request fails.

111 requests.ConnectionError

112 If network issue occurs.

113 ValueError

114 If response parsing fails.

115

116 """

117 if not time_series_codes and selection_pk is None:

118 raise ValueError(

119 "Either time_series_codes or selection_pk must be provided"

120 )

121

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

123

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

125 "respFormat": self.respformat,

126 "freqs": freq,

127 "startYear": start_year,

128 "endYear": end_year,

129 "isTidy": self.is_tidy,

130 "analysis": analysis,

131 }

132

133 if selection_pk is not None:

134 query_params["selectionPk"] = selection_pk

135 log_key = str(selection_pk)

136 else:

137 query_params["timeSeriesCodes"] = time_series_codes

138 log_key = time_series_codes

139

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

141

142 try:

143 response = requests.get(

144 url, params={**query_params, "auth_token": self.apikey}

145 )

146 response.raise_for_status()

147 except requests.ConnectionError as e:

148 raise requests.ConnectionError(

149 "Network error: Unable to connect to API"

150 ) from e

151 except requests.HTTPError as e:

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

153

154 if self.respformat == "csv":

155 try:

156 out: pd.DataFrame = (

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

158 )

159 except pd.errors.ParserError as e:

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

161 else:

162 try:

163 out: dict = response.json()

164 except ValueError as e:

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

166

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

168 return out

169

170 def _setup_cache(self) -> None:

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

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

173

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

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

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

177 if debug:

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

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

180

181 def _load_from_cache(

182 self, cache_key: str, resp_format: str

183 ) -> Optional[pd.DataFrame]:

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

185 if not self.use_cache:

186 return None

187

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

189 if not cache_path.exists():

190 return None

191

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

193 if resp_format == "parquet":

194 return pd.read_parquet(cache_path)

195 elif resp_format == "csv":

196 return pd.read_csv(cache_path)

197 return None

198

199 def _save_to_cache(

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

201 ) -> None:

202 """Save data to cache."""

203 if not self.use_cache:

204 return

205

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

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

208 if resp_format == "parquet":

209 data.to_parquet(cache_path, index=False)

210 elif resp_format == "csv":

211 data.to_csv(cache_path, index=False)

212

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

214 """

215 Fetch available recipes from Quantec API.

216

217 Returns

218 -------

219 Union[pd.DataFrame, dict]

220 DataFrame for CSV, dict for JSON.

221

222 Raises

223 ------

224 requests.HTTPError

225 If API request fails.

226 requests.ConnectionError

227 If network issue occurs.

228 ValueError

229 If response parsing fails.

230

231 """

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

233

234 try:

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

236 response.raise_for_status()

237 except requests.ConnectionError as e:

238 raise requests.ConnectionError(

239 "Network error: Unable to connect to API"

240 ) from e

241 except requests.HTTPError as e:

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

243

244 if self.respformat == "csv":

245 try:

246 recipes_data = response.json()

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

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

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

250 else:

251 try:

252 out: dict = response.json()

253 except ValueError as e:

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

255

256 log.debug(

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

258 )

259 return out

260

261 def get_selections(

262 self,

263 status: Optional[str] = None,

264 show: Optional[str] = None,

265 filter: Optional[str] = None,

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

267 """

268 Fetch user's available selections from Quantec API.

269

270 Parameters

271 ----------

272 status : Optional[str], optional

273 Filter by selection status using combined flags:

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

275 show : Optional[str], optional

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

277 filter : Optional[str], optional

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

279

280 Returns

281 -------

282 Union[pd.DataFrame, dict]

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

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

285

286 Raises

287 ------

288 requests.HTTPError

289 If API request fails.

290 requests.ConnectionError

291 If network issue occurs.

292 ValueError

293 If response parsing fails.

294

295 """

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

297

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

299

300 if status:

301 query_params["status"] = status

302 if show:

303 query_params["show"] = show

304 if filter:

305 query_params["filter"] = filter

306

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

308

309 try:

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

311 response.raise_for_status()

312 except requests.ConnectionError as e:

313 raise requests.ConnectionError(

314 "Network error: Unable to connect to API"

315 ) from e

316 except requests.HTTPError as e:

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

318

319 try:

320 resp = response.json()

321 if not resp:

322 selections_data = []

323 else:

324 # Transform data following the original logic

325 selections_data = [

326 {

327 "item": i,

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

329 "title": item["title"],

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

331 "is_owner": item["is_owner"],

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

333 "status": item["status"],

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

335 "modified": item["modified"],

336 }

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

338 ]

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

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

341

342 if self.respformat == "csv":

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

344 else:

345 out: dict = {"selections": selections_data}

346

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

348 return out

349

350 def get_grid_data(

351 self,

352 recipe_pk: int,

353 is_expanded: bool = True,

354 is_melted: bool = True,

355 resp_format: str = "dataframe",

356 selectdimensionnodes: dict = None,

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

358 """

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

360

361 Parameters

362 ----------

363 recipe_pk : int

364 Recipe primary key identifier.

365 is_expanded : bool, optional

366 Return expanded data format. Defaults to True.

367 is_melted : bool, optional

368 Return melted data format. Defaults to True.

369 resp_format : str, optional

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

371 selectdimensionnodes : dict, optional

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

373 Defaults to None.

374

375 Returns

376 -------

377 Union[pd.DataFrame, str, bytes]

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

379 or bytes if resp_format='parquet'.

380

381 Raises

382 ------

383 ValueError

384 If resp_format is invalid.

385 requests.HTTPError

386 If API request fails.

387 requests.ConnectionError

388 If network issue occurs.

389 ValueError

390 If response parsing fails.

391

392 """

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

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

395

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

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

398

399 # Check cache first

400 cache_key = self._generate_cache_key(

401 recipe_pk, is_expanded, is_melted, api_format, selectdimensionnodes

402 )

403

404 # Check if cached file exists and load raw data

405 if self.use_cache:

406 from pathlib import Path

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

408 if cache_path.exists():

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

410

411 if resp_format == "csv":

412 # Return cached CSV string

413 return cache_path.read_text()

414 elif resp_format == "parquet":

415 # Return cached parquet bytes

416 return cache_path.read_bytes()

417 else: # resp_format == "dataframe"

418 # Parse cached file to DataFrame

419 if api_format == "parquet":

420 cached_df = pd.read_parquet(cache_path)

421 else: # api_format == "csv"

422 cached_df = pd.read_csv(cache_path)

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

424

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

426

427 # Use POST if filtering, GET otherwise

428 if selectdimensionnodes:

429 # POST request for filtering

430 request_data = {

431 "respFormat": api_format,

432 "isExpanded": is_expanded,

433 "isMelted": is_melted,

434 "selectdimensionnodes": selectdimensionnodes,

435 }

436

437 headers = {

438 "Authorization": f"Token {self.apikey}",

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

440 }

441

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

443

444 try:

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

446 response.raise_for_status()

447 except requests.ConnectionError as e:

448 raise requests.ConnectionError(

449 "Network error: Unable to connect to API"

450 ) from e

451 except requests.HTTPError as e:

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

453 else:

454 # GET request (existing code)

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

456 "respFormat": api_format,

457 "isExpanded": is_expanded,

458 "isMelted": is_melted,

459 "auth_token": self.apikey,

460 }

461

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

463

464 try:

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

466 response.raise_for_status()

467 except requests.ConnectionError as e:

468 raise requests.ConnectionError(

469 "Network error: Unable to connect to API"

470 ) from e

471 except requests.HTTPError as e:

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

473

474 # Save raw response to cache first

475 if self.use_cache:

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

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

478 if api_format == "parquet":

479 cache_path.write_bytes(response.content)

480 else: # api_format == "csv"

481 cache_path.write_text(response.text)

482

483 # Handle return format based on user's request

484 if resp_format == "csv":

485 # Return raw CSV string

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

487 return response.text

488 elif resp_format == "parquet":

489 # Return raw parquet bytes

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

491 return response.content

492 else: # resp_format == "dataframe"

493 # Parse into DataFrame and apply cleaning

494 if api_format == "parquet":

495 try:

496 from io import BytesIO

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

498 except Exception as e:

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

500 else: # api_format == "csv"

501 try:

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

503 except pd.errors.ParserError as e:

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

505

506 # Clean up data (only for DataFrame output)

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

508

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

510 return out