Add support for EIA forecasts endpoint

docs/endpoints/about-endpoints.md

@@ -10,7 +10,8 @@ The VortexaSDK currently contains the following endpoints:
 4. Geographies
 5. Products
 6. Vessels
-
+7. Cargo Time Series
+8. EIA Forecasts
 
 Each endpoint offers either one, or both, of two different functionalities:
 

pydocmd.yml

@@ -38,6 +38,10 @@ generate:
       - vortexasdk.api.cargo_movement++
   - entities/vessel_movement.md:
       - vortexasdk.api.vessel_movement++
+  - endpoints/eia_forecasts.md:
+      - vortexasdk.endpoints.eia_forecasts++
+      - vortexasdk.endpoints.eia_forecasts_result++
+
 pages:
   - Home: index.md << README.md
   - Endpoints:
@@ -50,6 +54,7 @@ pages:
       Geographies: endpoints/geographies.md
       Products: endpoints/products.md
       Vessels: endpoints/vessels.md
+      EIA Forecasts: endpoints/eia_forecasts.md
   - Entities:
       Cargo Movement: entities/cargo_movement.md
       Corporation: entities/corporation.md

tests/endpoints/test_eia_forecasts_real.py

@@ -0,0 +1,47 @@
+from tests.testcases import TestCaseUsingRealAPI
+from datetime import datetime
+from tests.timer import Timer
+from vortexasdk.endpoints.eia_forecasts import EIAForecasts
+from docs.utils import to_markdown
+
+
+class TestEIAForecastsReal(TestCaseUsingRealAPI):
+    def test_search_preset_crude_imports(self):
+
+        preset = "padd1-crude-imports"
+        filter_time_min = datetime(2020, 1, 20)
+        filter_time_max = datetime(2020, 1, 24)
+
+        forecasts = EIAForecasts().search(
+            preset=preset,
+            filter_time_min=filter_time_min,
+            filter_time_max=filter_time_max,
+        )
+        assert len(forecasts) == 1
+        values = [g.date for g in forecasts.to_list()]
+
+        assert "2020-01-24T00:00:00.000Z" in values
+
+        print(to_markdown(forecasts.to_df()))
+
+    def test_search_preset_gasoline_exports(self):
+
+        preset = "us-gasoline-exports"
+        filter_time_min = datetime(2020, 3, 20)
+        filter_time_max = datetime(2020, 4, 20)
+
+        forecasts = EIAForecasts().search(
+            preset=preset,
+            filter_time_min=filter_time_min,
+            filter_time_max=filter_time_max,
+        )
+        assert len(forecasts) == 5
+
+        values = [g.date for g in forecasts.to_list()]
+        assert "2020-04-17T00:00:00.000Z" in values
+        assert "2020-04-10T00:00:00.000Z" in values
+        assert "2020-04-03T00:00:00.000Z" in values
+        assert "2020-03-27T00:00:00.000Z" in values
+        assert "2020-03-20T00:00:00.000Z" in values
+
+        print(to_markdown(forecasts.to_df()))

vortexasdk/__init__.py

@@ -9,6 +9,7 @@
     Products,
     VesselMovements,
     Vessels,
+    EIAForecasts,
 )
 
 # noinspection PyUnresolvedReferences

vortexasdk/api/__init__.py

@@ -19,3 +19,4 @@
 )
 from vortexasdk.api.vessel import Vessel, VesselEntity
 from vortexasdk.api.vessel_movement import VesselEvent, VesselMovement
+from vortexasdk.api.eia_forecast import EIAForecast

vortexasdk/api/eia_forecast.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from vortexasdk.api.serdes import FromDictMixin
+
+
+@dataclass(frozen=True)
+class EIAForecast(FromDictMixin):
+    """Represent a EIA forecast record returned by the API."""
+
+    date: str
+    forecast_fri: float
+    value: Optional[int]
+    stocks: Optional[int]
+    cover: Optional[float]
+    runs: Optional[float]

vortexasdk/endpoints/__init__.py

@@ -7,3 +7,4 @@
 from vortexasdk.endpoints.geographies import Geographies
 from vortexasdk.endpoints.products import Products
 from vortexasdk.endpoints.vessels import Vessels
+from vortexasdk.endpoints.eia_forecasts import EIAForecasts

vortexasdk/endpoints/eia_forecasts.py

@@ -0,0 +1,88 @@
+"""EIA Forecasts Endpoint."""
+
+from vortexasdk.endpoints.eia_forecasts_result import EIAForecastResult
+from vortexasdk.endpoints.endpoints import EIA_FORECASTS_RESOURCE
+from vortexasdk.operations import Search
+from vortexasdk.api.shared_types import to_ISODate
+from datetime import datetime
+
+
+class EIAForecasts(Search):
+    """EIA forecasts Endpoint, use this to search through Vortexa's EIA Forecasts data.
+
+    The data includes:
+
+    - `date`: date of the forecast
+    - `forecast_fri`: Vortexa's data science based forecast of the EIA number to be published on the week
+    - `value`: Actual EIA import/export numbers as published by the EIA Weekly Supply Estimates report
+    - `stocks`: EIA stocks (kbl)
+    - `cover`: Cover (days of Supply for the whole of the US, as published by the EIA Weekly Supply Estimates report)
+    - `runs`: refinery runs (refiner “Percent Operable Utilization” as published by the EIA Weekly Supply Estimates report)
+
+    """
+
+    def __init__(self):
+        Search.__init__(self, EIA_FORECASTS_RESOURCE)
+
+    def search(
+        self,
+        preset: str = "padd1-gasoline-imports",
+        filter_time_min: datetime = datetime(2020, 1, 1),
+        filter_time_max: datetime = datetime(2020, 1, 31),
+    ) -> EIAForecastResult:
+        """
+        Find EIA forecasts for a given preset and date range.
+
+        # Arguments
+            preset: Use to specify what geography and product information you would like to query.
+            Preset can be: 'padd1-gasoline-imports', 'padd3-gasoline-imports', 'padd5-gasoline-imports',
+            'us-gasoline-exports', 'padd1-crude-imports', 'padd3-crude-imports', 'padd5-crude-imports',
+            'us-crude-exports', 'padd1-diesel-imports', 'padd3-diesel-imports', 'padd5-diesel-imports',
+            'us-diesel-exports', 'padd1-jet-imports', 'padd5-jet-imports', 'us-jet-exports',
+            'padd1-fueloil-imports', 'padd3-fueloil-imports', 'padd5-fueloil-imports' or 'us-fueloil-exports'
+
+            filter_time_min: The UTC start date of the time filter
+
+            filter_time_max: The UTC end date of the time filter
+
+        # Returns
+        List of EIA Forecast object matching selected 'preset'.
+
+        # Examples
+
+        Find PADD5 gasoline imports EIA forecasts from January 2019.
+        ```python
+        >>> from datetime import datetime
+        >>> from vortexasdk import EIAForecasts
+        >>> df = EIAForecasts().search(
+        ...     preset="padd5-gasoline-imports",
+        ...     filter_time_min=datetime(2020, 1, 1),
+        ...     filter_time_max=datetime(2020, 1, 31)
+        ... ).to_df()
+
+        ```
+
+        returns
+
+        | date                     | forecast_fri     | value | stocks | cover | runs |
+        | ------------------------ | ---------------- | ----- | ------ | ----- | ---- |
+        | 2020-01-31T00:00:00.000Z | 454.96048964485  | 323   | 9541   | 26.5  | 65.9 |
+        | 2020-01-24T00:00:00.000Z | 545.453497230504 | 579   | 10461  | 25.9  | 61.5 |
+        | 2020-01-17T00:00:00.000Z | 510.289752707662 | 549   | 10325  | 25.2  | 64.7 |
+        | 2020-01-10T00:00:00.000Z | 469.841470826967 |       |        |       |      |
+        | 2020-01-03T00:00:00.000Z | 640.443229654771 |       |        |       |      |
+
+        Some values can be NULL: value, stocks, cover, runs. It can happen when:
+
+        - it's a very recent forecast, the Vortexa's data science based forecast (forecast_fri) is available but
+          the complete EIA data isn't yet
+        - it's an older forecast and the data is not available
+
+        """
+        search_params = {
+            "preset": preset,
+            "filter_time_min": to_ISODate(filter_time_min),
+            "filter_time_max": to_ISODate(filter_time_max),
+        }
+
+        return EIAForecastResult(super().search(**search_params))

vortexasdk/endpoints/eia_forecasts_result.py

@@ -0,0 +1,41 @@
+from typing import List
+
+import pandas as pd
+
+from vortexasdk.api import EIAForecast
+from vortexasdk.api.search_result import Result
+from vortexasdk.logger import get_logger
+from vortexasdk.result_conversions import create_dataframe, create_list
+
+logger = get_logger(__name__)
+
+
+class EIAForecastResult(Result):
+    """Container class that holds the result obtained from calling the `EIAForecasts` endpoint."""
+
+    def to_list(self) -> List[EIAForecast]:
+        """Represent EIAForecast data as a list."""
+        # noinspection PyTypeChecker
+        return create_list(super().to_list(), EIAForecast)
+
+    def to_df(self, columns=None) -> pd.DataFrame:
+        """
+        Represent EIA forecasts as a `pd.DataFrame`.
+
+        # Arguments
+            columns: The EIA forecasts columns we want in the dataframe. Enter `columns='all'` to include all columns.
+            Defaults to `columns = ['date', 'forecast_fri', 'value', 'stocks', 'cover', 'runs']`.
+
+
+        # Returns
+        `pd.DataFrame` of EIA forecasts.
+        """
+        return create_dataframe(
+            columns=columns,
+            default_columns=DEFAULT_COLUMNS,
+            data=super().to_list(),
+            logger_description="EIAForecasts",
+        )
+
+
+DEFAULT_COLUMNS = ["date", "forecast_fri", "value", "stocks", "cover", "runs"]

vortexasdk/endpoints/endpoints.py

@@ -2,6 +2,7 @@
 
 CARGO_MOVEMENTS_RESOURCE = "/cargo-movements/search"
 VESSEL_MOVEMENTS_RESOURCE = "/vessel-movements/search"
+EIA_FORECASTS_RESOURCE = "/search/eia-forecasts"
 
 ATTRIBUTES_REFERENCE = "/reference/attributes"
 GEOGRAPHIES_REFERENCE = "/reference/geographies"