From 07989ca2608832825d5b6beb2f7f266fcf363eb4 Mon Sep 17 00:00:00 2001
From: philsv <frphsv@gmail.com>
Date: Thu, 26 Jan 2023 12:49:42 +0100
Subject: [PATCH] feat: add simpler api call by series_id only

---
 README.md           |  9 ++++++--
 myeia/api.py        | 56 +++++++++++++++++++++++++++++++++++++--------
 myeia/version.py    |  2 +-
 tests/test_myeia.py | 18 +++++++++++----
 4 files changed, 69 insertions(+), 16 deletions(-)

diff --git a/README.md b/README.md
index 0d01104..ca29d9e 100644
--- a/README.md
+++ b/README.md
@@ -40,9 +40,12 @@ EIA_TOKEN=YOUR_TOKEN_HERE
 ```
 
 Lets look at an example of how to get the *EIA Natural Gas Futures*.
+You can use the simpler v1 API method where you only need to pass the `series_id` or you can use the newer v2 API method where you need to pass the `route`, `series`, and `frequency`.
 
 ```python
-df = eia.get_data(
+df = eia.get_series(series_id="NG.RNGC1.D")
+
+df = eia.get_series_via_route(
     route="natural-gas/pri/fut",
     series="RNGC1",
     frequency="daily",
@@ -66,7 +69,9 @@ Date
 Lets look at another example the *Total OPEC Petroleum Supply* where the facet is available as `seriesId`. By Default it is set as `series` but we can define the facet as `seriesId`.
 
 ```python
-df = eia.get_data(
+df = eia.get_series(series_id="STEO.PAPR_OPEC.M")
+
+df = eia.get_series_via_route(
     route="steo",
     series="PAPR_OPEC",
     frequency="monthly",
diff --git a/myeia/api.py b/myeia/api.py
index 67672e0..c761a17 100644
--- a/myeia/api.py
+++ b/myeia/api.py
@@ -20,7 +20,42 @@ class API:
     token: Optional[str] = os.getenv("EIA_TOKEN")
     url: str = "https://api.eia.gov/v2/"
 
-    def get_data(
+    def get_series(
+        self,
+        series_id: str,
+        new_name: str = "",
+    ) -> pd.DataFrame:
+        """
+        Returns data for a given series in the simpler APIv1 format.
+
+        Args:
+            series_id (str): The series ID. For example, "NG.RNGC1.W".
+        Returns:
+            pd.DataFrame: A DataFrame with the date and value columns.
+        """
+        headers = {"Accept": "*/*"}
+        url = f"{self.url}seriesid/{series_id}?api_key={self.token}"
+        response = requests.get(url, headers=headers)
+        response.raise_for_status()
+        json_response = response.json()
+
+        base_df = pd.DataFrame(json_response["response"]["data"])
+
+        if "series-description" in base_df.columns:
+            series_description = base_df["series-description"][0]
+        else:
+            series_description = series_id
+
+        if new_name != "":
+            series_description = new_name
+
+        df = base_df[["period", "value"]]
+        df.rename(columns={df.columns[0]: "Date", df.columns[1]: series_description}, inplace=True)
+        df["Date"] = pd.to_datetime(df["Date"])
+        df.set_index("Date", inplace=True)
+        return df
+
+    def get_series_via_route(
         self,
         route: str,
         series: str,
@@ -29,19 +64,22 @@ def get_data(
         new_name: str = "",
     ) -> pd.DataFrame:
         """
-        Returns data for a given series.
+        Returns data for a given series in the newer APIv2 format.
+
+        Args:
+            route (str): The route to the series. For example, "natural-gas/pri/fut".
+            series (str): The series ID. For example, "RNGC1".
+            frequency (str): The frequency of the series. For example, "daily".
+            facet (str): The facet of the series. For example, "series", "seriesId".
+            new_name (str): A name you want to give the series.
+        Returns:
+            pd.DataFrame: A DataFrame with the date and value columns.
         """
-
-        headers = {
-            "Accept": "*/*",
-        }
+        headers = {"Accept": "*/*"}
 
         api_route = f"{route}/data/?api_key={self.token}&data[]=value&frequency={frequency}"
-
         series = f"&facets[{facet}][]={series}"
-
         sort = "&sort[0][column]=period&sort[0][direction]=desc"
-
         url = f"{self.url}{api_route}{series}{sort}"
 
         response = requests.get(url, headers=headers)
diff --git a/myeia/version.py b/myeia/version.py
index b5fdc75..d31c31e 100644
--- a/myeia/version.py
+++ b/myeia/version.py
@@ -1 +1 @@
-__version__ = "0.2.2"
+__version__ = "0.2.3"
diff --git a/tests/test_myeia.py b/tests/test_myeia.py
index 3e484e6..89b7c3e 100644
--- a/tests/test_myeia.py
+++ b/tests/test_myeia.py
@@ -3,9 +3,19 @@
 
 from myeia.api import API
 
+eia = API()
 
-@pytest.mark.parametrize("route,series,frequency,facet", [("natural-gas/pri/fut", "RNGC1", "daily", "series"), ("petroleum/move/pipe", "MD0MP_R10-R20_1", "monthly", "series")])
-def test_get_data(route, series, frequency, facet):
-    eia = API()
-    df = eia.get_data(route, series, frequency, facet)
+
+@pytest.mark.parametrize("series_id", ["NG.RNGC1.D", "PET.WCESTUS1.W", "PET.MD0MP_R10-R20_1.M", "INTL.29-12-HKG-BKWH.A"])
+def test_get_series(series_id):
+    df = eia.get_series(series_id)
+    assert isinstance(df, pd.DataFrame)
+
+
+@pytest.mark.parametrize(
+    "route,series,frequency,facet",
+    [("natural-gas/pri/fut", "RNGC1", "daily", "series"), ("petroleum/stoc/wstk", "WCESTUS1", "weekly", "series"), ("petroleum/move/pipe", "MD0MP_R10-R20_1", "monthly", "series")],
+)
+def test_get_series_via_route(route, series, frequency, facet):
+    df = eia.get_series_via_route(route, series, frequency, facet)
     assert isinstance(df, pd.DataFrame)