Add support for EIA forecasts endpoint

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,54 @@
+from tests.testcases import TestCaseUsingRealAPI
+from datetime import datetime
+from tests.timer import Timer
+from vortexasdk.endpoints.eia_forecasts import EIAForecasts
+
+
+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,
+            )
+            .to_list()
+        )
+        assert len(forecasts) == 1
+
+        values = [g.name for g in forecasts]
+        assert "2020-01-24T00:00:00.000Z" in values
+
+        print(to_markdown(forecasts))
+
+    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,
+            )
+            .to_list()
+        )
+        assert len(forecasts) == 5
+
+        values = [g.name for g in forecasts]
+        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))

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,17 @@
+from dataclasses import dataclass
+from typing import List
+
+from vortexasdk.api.serdes import FromDictMixin
+from vortexasdk.api.shared_types import IDName
+
+
+@dataclass(frozen=True)
+class EIAForecast(IDName, FromDictMixin):
+    """Represent a EIA forecast record returned by the API."""
+
+    date: List[str]
+    forecast_fri: float
+    value: int
+    stocks: int
+    cover: float
+    runs: 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,68 @@
+"""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
+
+
+class EIAForecasts(Search):
+    """EIA forecasts Endpoint."""
+
+    def __init__(self):
+        Search.__init__(self, EIA_FORECASTS_RESOURCE)
+
+    def search(
+        self,
+        preset: str = None,
+        filter_time_min: str = False,
+        filter_time_max: str = False,
+    ) -> EIAForecastResult:
+        """
+        Find EIA forecasts for a given date preset and date range.
+
+        # Arguments
+            preset: the EIA forecasts preset to be returned. 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 'preset'
+
+        # Examples
+
+        Find PADD5 gasoline imports EIA forecasts from January 2019.
+        ```python
+        >>> from vortexasdk import EIAForecasts
+        >>> df = EIAForecasts().search(
+                type="padd5-gasoline-imports",
+                filter_time_min="2020-01-01T00:00:00.000Z",
+                filter_time_max="2020-01-31T00:00:00.000Z"
+            ).to_df()
+        ```
+
+        returns
+
+        | date                     | forecast_fri     | value | stocks | cover | runs |
+        | ------------------------ | ---------------- | ----- | ------ | ----- | ---- |
+        | 2019-01-25T00:00:00.000Z | 48.6277718725167 | 0     | 32811  | 28.8  | 89.4 |
+        | 2019-01-18T00:00:00.000Z | 29.5812233704497 | 48    | 32426  | 29.8  | 88.7 |
+        | 2019-01-11T00:00:00.000Z | 45.5004976086444 | 0     | 31471  | 29    | 91.8 |
+        | 2019-01-04T00:00:00.000Z | 0                |       |        |       |      |
+
+
+        """
+        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 features we want in the dataframe. Enter `columns='all'` to include all features.
+            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"