feat: basic api keys #23

Author: dshemetov

epidatpy/__init__.py

@@ -22,5 +22,6 @@
     GeoType,
     TimeType,
 )
+from ._auth import get_api_key
 
 __author__ = "Delphi Group"

epidatpy/_auth.py

@@ -0,0 +1,16 @@
+import os
+import warnings
+from typing import Optional
+
+
+def get_api_key() -> Optional[str]:
+    key = os.environ.get("DELPHI_EPIDATA_KEY", None)
+
+    if not key:
+        warnings.warn(
+            "DELPHI_EPIDATA_KEY environment variable not set. "
+            "Please set it to your Epidata API key to avoid rate limits. "
+            "You can get a free key at: https://api.delphi.cmu.edu/epidata/admin/registration_form"
+        )
+
+    return key

epidatpy/request.py

@@ -3,6 +3,7 @@
 from json import loads
 
 from requests import Response, Session
+from requests.auth import HTTPBasicAuth
 from tenacity import retry, stop_after_attempt
 from pandas import DataFrame
 
@@ -19,18 +20,20 @@
 from ._endpoints import AEpiDataEndpoints
 from ._constants import HTTP_HEADERS, BASE_URL
 from ._covidcast import CovidcastDataSources, define_covidcast_fields
+from ._auth import get_api_key
 
 
 @retry(reraise=True, stop=stop_after_attempt(2))
 def _request_with_retry(
     url: str, params: Mapping[str, str], session: Optional[Session] = None, stream: bool = False
 ) -> Response:
     """Make request with a retry if an exception is thrown."""
+    basic_auth = HTTPBasicAuth("epidata", get_api_key())
 
     def call_impl(s: Session) -> Response:
-        res = s.get(url, params=params, headers=HTTP_HEADERS, stream=stream)
+        res = s.get(url, params=params, headers=HTTP_HEADERS, stream=stream, auth=basic_auth)
         if res.status_code == 414:
-            return s.post(url, params=params, headers=HTTP_HEADERS, stream=stream)
+            return s.post(url, params=params, headers=HTTP_HEADERS, stream=stream, auth=basic_auth)
         return res
 
     if session:

tests/test_auth.py

@@ -0,0 +1,12 @@
+from pytest import warns, MonkeyPatch
+
+from epidatpy import get_api_key
+
+
+def test_get_api_key(monkeypatch: MonkeyPatch) -> None:
+    with monkeypatch.context() as m:
+        m.setenv("DELPHI_EPIDATA_KEY", "test")
+        assert get_api_key() == "test"
+        m.delenv("DELPHI_EPIDATA_KEY")
+        with warns(UserWarning):
+            assert get_api_key() is None