Refactor API class and add type hints

myeia/api.py

@@ -17,6 +17,7 @@
 
 load_dotenv(verbose=True)
 
+
 class API:
     """
     Python Wrapper for U.S. Energy Information Administration (EIA) APIv2.
@@ -32,9 +33,11 @@ def __init__(
         if token:
             self.token = token
         elif os.getenv("EIA_TOKEN"):
-            self.token = os.getenv("EIA_TOKEN")
+            self.token = str(os.getenv("EIA_TOKEN"))
         else:
-             raise ValueError('EIA_TOKEN is not set. Please set it in your .env file or environment variables.')
+            raise ValueError(
+                "EIA_TOKEN is not set. Please set it in your .env file or environment variables."
+            )
 
         self.base_url = "https://api.eia.gov/v2/"
         self.header = {"Accept": "*/*"}

requirements.txt

@@ -1,8 +1,8 @@
 backoff>=2.2.1
+numpy>=1.21.4
 pandas>=1.5.3
-pytest>=7.2.1
 pytest-mock>=3.14.0
+pytest>=7.2.1
 python-dateutil>=2.9.0
 python-dotenv>=0.19.0
-requests>=2.32.0
-numpy>=1.21.4
+requests>=2.32.0

setup.py

@@ -1,4 +1,4 @@
-import setuptools
+import setuptools  # type: ignore
 
 from myeia.version import __version__
 

tests/test_api.py

@@ -9,8 +9,13 @@
 
 eia = API()
 
-def mock_requests_get(mocker, file_path):
+
+def mock_requests_get(
+    mocker,
+    file_path: str,
+) -> requests.Response:
     """Helper function to mock requests.get and manage test data files."""
+
     class MockGetResponse:
         def status_code(self):
             return 200
@@ -31,10 +36,10 @@ def raise_for_status(self):
 
     # register spy on requests.get
     spy_get = mocker.spy(requests, "get")
-
     return spy_get
 
-def save_mock_data(file_path, json_response):
+
+def save_mock_data(file_path: str, json_response: dict) -> None:
     """Helper function to save mock data to a file."""
     # clean request data as it can contain secrets like api key
     json_response["request"] = {}
@@ -44,12 +49,14 @@ def save_mock_data(file_path, json_response):
     with open(file_path, "w") as f:
         f.write(json.dumps(json_response))
 
-def get_mock_data_path(file_path):
+
+def get_mock_data_path(file_path: str) -> str:
     """Helper function to get the path of the mock data file."""
     base_dir = os.path.dirname(os.path.abspath(__file__))
     data_dir = os.path.join(base_dir, "data/")
     return os.path.join(data_dir, file_path)
 
+
 @pytest.mark.parametrize(
     "series_id, start_date, end_date",
     [
@@ -62,7 +69,6 @@ def get_mock_data_path(file_path):
 def test_get_series(series_id, start_date, end_date, mocker):
     """Test get_series method."""
     file_path = get_mock_data_path(f"{series_id}_{start_date}_{end_date}.json")
-
     spy_get = mock_requests_get(mocker, file_path)
 
     df = eia.get_series(series_id, start_date=start_date, end_date=end_date)
@@ -75,6 +81,7 @@ def test_get_series(series_id, start_date, end_date, mocker):
     assert not df.empty
     assert isinstance(df, pd.DataFrame)
 
+
 @pytest.mark.parametrize(
     "route, series, frequency, facet",
     [
@@ -88,8 +95,9 @@ def test_get_series(series_id, start_date, end_date, mocker):
 )
 def test_get_series_via_route(route, series, frequency, facet, mocker):
     """Test get_series_via_route method."""
-    file_path = get_mock_data_path(f"{route.replace('/', '-')}_{series}_{frequency}_{facet}.json")
-
+    file_path = get_mock_data_path(
+        f"{route.replace('/', '-')}_{series}_{frequency}_{facet}.json"
+    )
     spy_get = mock_requests_get(mocker, file_path)
 
     df = eia.get_series_via_route(route, series, frequency, facet)
@@ -100,4 +108,4 @@ def test_get_series_via_route(route, series, frequency, facet, mocker):
         save_mock_data(file_path, spy_get.spy_return.json())
 
     assert not df.empty
-    assert isinstance(df, pd.DataFrame)
+    assert isinstance(df, pd.DataFrame)