Caching

Author: rzats

epidatpy/_model.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
 from datetime import date
 from enum import Enum
+from os import environ
 from typing import (
     Final,
     List,
@@ -146,6 +147,7 @@ class AEpiDataCall:
     meta: Final[Sequence[EpidataFieldInfo]]
     meta_by_name: Final[Mapping[str, EpidataFieldInfo]]
     only_supports_classic: Final[bool]
+    use_cache: Final[bool]
 
     def __init__(
         self,
@@ -154,13 +156,18 @@ def __init__(
         params: Mapping[str, Optional[EpiRangeParam]],
         meta: Optional[Sequence[EpidataFieldInfo]] = None,
         only_supports_classic: bool = False,
+        use_cache: Optional[bool] = None,
     ) -> None:
         self._base_url = base_url
         self._endpoint = endpoint
         self._params = params
         self.only_supports_classic = only_supports_classic
         self.meta = meta or []
         self.meta_by_name = {k.name: k for k in self.meta}
+        # Set the use_cache value from the constructor if present.
+        # Otherwise check the USE_EPIDATPY_CACHE variable, accepting various "truthy" values.
+        self.use_cache = use_cache \
+            or (environ.get("USE_EPIDATPY_CACHE", "").lower() in ['true', 't', '1'])
 
     def _verify_parameters(self) -> None:
         # hook for verifying parameters before sending

epidatpy/request.py

@@ -10,6 +10,9 @@
     cast,
 )
 
+from appdirs import user_cache_dir
+from diskcache import Cache
+from json import dumps
 from pandas import CategoricalDtype, DataFrame, Series, to_datetime
 from requests import Response, Session
 from requests.auth import HTTPBasicAuth
@@ -33,7 +36,7 @@
 
 # Make the linter happy about the unused variables
 __all__ = ["Epidata", "EpiDataCall", "EpiDataContext", "EpiRange", "CovidcastEpidata"]
-
+CACHE_DIRECTORY = user_cache_dir(appname="epidatpy", appauthor="delphi")
 
 @retry(reraise=True, stop=stop_after_attempt(2))
 def _request_with_retry(
@@ -73,8 +76,9 @@ def __init__(
         params: Mapping[str, Optional[EpiRangeParam]],
         meta: Optional[Sequence[EpidataFieldInfo]] = None,
         only_supports_classic: bool = False,
+        use_cache = None,
     ) -> None:
-        super().__init__(base_url, endpoint, params, meta, only_supports_classic)
+        super().__init__(base_url, endpoint, params, meta, only_supports_classic, use_cache)
         self._session = session
 
     def with_base_url(self, base_url: str) -> "EpiDataCall":
@@ -100,13 +104,23 @@ def classic(
         """Request and parse epidata in CLASSIC message format."""
         self._verify_parameters()
         try:
+            if self.use_cache:
+                with Cache(CACHE_DIRECTORY) as cache:
+                    cache_key = str(self._endpoint) + str(self._params)
+                    if cache_key in cache:
+                        return cache[cache_key]
             response = self._call(fields)
             r = cast(EpiDataResponse, response.json())
             if disable_type_parsing:
                 return r
             epidata = r.get("epidata")
             if epidata and isinstance(epidata, list) and len(epidata) > 0 and isinstance(epidata[0], dict):
                 r["epidata"] = [self._parse_row(row, disable_date_parsing=disable_date_parsing) for row in epidata]
+            if self.use_cache:
+                with Cache(CACHE_DIRECTORY) as cache:
+                    cache_key = str(self._endpoint) + str(self._params)
+                    # Set TTL to 7 days (TODO: configurable?)
+                    cache.set(cache_key, r, expire=7*24*60*60)
             return r
         except Exception as e:  # pylint: disable=broad-except
             return {"result": 0, "message": f"error: {e}", "epidata": []}
@@ -130,6 +144,13 @@ def df(
         if self.only_supports_classic:
             raise OnlySupportsClassicFormatException()
         self._verify_parameters()
+
+        if self.use_cache:
+            with Cache(CACHE_DIRECTORY) as cache:
+                cache_key = str(self._endpoint) + str(self._params)
+                if cache_key in cache:
+                    return cache[cache_key]
+
         json = self.classic(fields, disable_type_parsing=True)
         rows = json.get("epidata", [])
         pred = fields_to_predicate(fields)
@@ -175,6 +196,13 @@ def df(
                     df[info.name] = to_datetime(df[info.name], format="%Y%m%d")
                 except ValueError:
                     pass
+
+        if self.use_cache:
+            with Cache(CACHE_DIRECTORY) as cache:
+                cache_key = str(self._endpoint) + str(self._params)
+                # Set TTL to 7 days (TODO: configurable?)
+                cache.set(cache_key, df, expire=7*24*60*60)
+
         return df
 
 
@@ -203,8 +231,9 @@ def _create_call(
         params: Mapping[str, Optional[EpiRangeParam]],
         meta: Optional[Sequence[EpidataFieldInfo]] = None,
         only_supports_classic: bool = False,
+        use_cache: bool = False,
     ) -> EpiDataCall:
-        return EpiDataCall(self._base_url, self._session, endpoint, params, meta, only_supports_classic)
+        return EpiDataCall(self._base_url, self._session, endpoint, params, meta, only_supports_classic, use_cache)
 
 
 Epidata = EpiDataContext()

pyproject.toml

@@ -31,6 +31,8 @@ classifiers = [
 requires-python = ">=3.8"
 dependencies = [
     "aiohttp",
+    "appdirs",
+    "diskcache",
     "epiweeks>=2.1",
     "pandas>=1",
     "requests>=2.25",