Set vat on class or function level with new enum (#271)

* Add VatCategory enum and update tests

* Update examples and documentation

* Rename VatCategory to VatOption

README.md

@@ -71,12 +71,12 @@ The gas prices do not change per hour, but are fixed for 24 hours. Which means t
 import asyncio
 
 from datetime import date
-from energyzero import EnergyZero
+from energyzero import EnergyZero, VatOption
 
 
 async def main() -> None:
     """Show example on fetching the energy prices from EnergyZero."""
-    async with EnergyZero(incl_vat=True) as client:
+    async with EnergyZero(vat=VatOption.INCLUDE) as client:
         start_date = date(2022, 12, 7)
         end_date = date(2022, 12, 7)
 
@@ -92,7 +92,7 @@ if __name__ == "__main__":
 
 | Parameter | value Type | Description |
 | :-------- | :--------- | :---------- |
-| `incl_vat` | bool (default: **True**) | Include or exclude VAT |
+| `vat` | enum (default: **VatOption.INCLUDE**) | Include or exclude VAT on class level |
 
 ### Function Parameters
 
@@ -101,6 +101,7 @@ if __name__ == "__main__":
 | `start_date` | datetime | The start date of the selected period |
 | `end_date` | datetime | The end date of the selected period |
 | `interval` | integer (default: **4**) | The interval of data return (**day**, **week**, **month**, **year**) |
+| `vat` | enum (default: class value) | Include or exclude VAT (**VatOption.INCLUDE** or **VatOption.EXCLUDE**) |
 
 **Interval**
 4: Dag

examples/energy.py

@@ -5,15 +5,15 @@
 
 import pytz
 
-from energyzero import EnergyZero
+from energyzero import EnergyZero, VatOption
 
 
 async def main() -> None:
     """Show example on fetching the energy prices from EnergyZero."""
-    async with EnergyZero() as client:
+    async with EnergyZero(vat=VatOption.INCLUDE) as client:
         local = pytz.timezone("CET")
-        today = date(2023, 8, 9)
-        tomorrow = date(2023, 8, 9)
+        today = date(2023, 12, 5)
+        tomorrow = date(2023, 12, 6)
 
         energy_today = await client.energy_prices(start_date=today, end_date=today)
         energy_tomorrow = await client.energy_prices(

examples/gas.py

@@ -3,13 +3,13 @@
 import asyncio
 from datetime import date, timedelta
 
-from energyzero import EnergyZero
+from energyzero import EnergyZero, VatOption
 
 
 async def main() -> None:
     """Show example on fetching the gas prices from EnergyZero."""
-    async with EnergyZero() as client:
-        today = date(2023, 3, 28)
+    async with EnergyZero(vat=VatOption.INCLUDE) as client:
+        today = date(2023, 12, 5)
 
         gas_today = await client.gas_prices(start_date=today, end_date=today)
         print()

src/energyzero/__init__.py

@@ -1,5 +1,6 @@
 """Asynchronous Python client for the EnergyZero API."""
 
+from .const import VatOption
 from .energyzero import EnergyZero
 from .exceptions import (
     EnergyZeroConnectionError,
@@ -9,10 +10,11 @@
 from .models import Electricity, Gas
 
 __all__ = [
-    "Gas",
     "Electricity",
     "EnergyZero",
+    "EnergyZeroConnectionError",
     "EnergyZeroError",
     "EnergyZeroNoDataError",
-    "EnergyZeroConnectionError",
+    "Gas",
+    "VatOption",
 ]

src/energyzero/const.py

@@ -0,0 +1,11 @@
+"""Constants for EnergyZero API client."""
+from __future__ import annotations
+
+from enum import Enum
+
+
+class VatOption(str, Enum):
+    """Enum representing whether to include VAT or not."""
+
+    INCLUDE = "true"
+    EXCLUDE = "false"

src/energyzero/energyzero.py

@@ -12,6 +12,7 @@
 from aiohttp.hdrs import METH_GET
 from yarl import URL
 
+from .const import VatOption
 from .exceptions import (
     EnergyZeroConnectionError,
     EnergyZeroError,
@@ -24,7 +25,7 @@
 class EnergyZero:
     """Main class for handling data fetching from EnergyZero."""
 
-    incl_vat: bool = True
+    vat: VatOption = VatOption.INCLUDE
     request_timeout: float = 10.0
     session: ClientSession | None = None
 
@@ -107,6 +108,7 @@ async def gas_prices(
         start_date: date,
         end_date: date,
         interval: int = 4,
+        vat: VatOption | None = None,
     ) -> Gas:
         """Get gas prices for a given period.
 
@@ -115,6 +117,7 @@ async def gas_prices(
             start_date: Start date of the period.
             end_date: End date of the period.
             interval: Interval of the prices.
+            vat: VAT category.
 
         Returns:
         -------
@@ -176,7 +179,7 @@ async def gas_prices(
                 "tillDate": utc_end_date.strftime("%Y-%m-%dT%H:%M:%S.999Z"),
                 "interval": interval,
                 "usageType": 3,
-                "inclBtw": "true" if self.incl_vat else "false",
+                "inclBtw": vat.value if vat is not None else self.vat.value,
             },
         )
 
@@ -190,6 +193,7 @@ async def energy_prices(
         start_date: date,
         end_date: date,
         interval: int = 4,
+        vat: VatOption | None = None,
     ) -> Electricity:
         """Get energy prices for a given period.
 
@@ -198,6 +202,7 @@ async def energy_prices(
             start_date: Start date of the period.
             end_date: End date of the period.
             interval: Interval of the prices.
+            vat: VAT category.
 
         Returns:
         -------
@@ -234,7 +239,7 @@ async def energy_prices(
                 "tillDate": utc_end_date.strftime("%Y-%m-%dT%H:%M:%S.999Z"),
                 "interval": interval,
                 "usageType": 1,
-                "inclBtw": "true" if self.incl_vat else "false",
+                "inclBtw": vat.value if vat is not None else self.vat.value,
             },
         )
 

tests/test_models.py

@@ -5,7 +5,7 @@
 from aiohttp import ClientSession
 from aresponses import ResponsesMockServer
 
-from energyzero import Electricity, EnergyZero, EnergyZeroNoDataError, Gas
+from energyzero import Electricity, EnergyZero, EnergyZeroNoDataError, Gas, VatOption
 
 from . import load_fixtures
 
@@ -29,6 +29,7 @@ async def test_electricity_model(aresponses: ResponsesMockServer) -> None:
         energy: Electricity = await client.energy_prices(
             start_date=today,
             end_date=today,
+            vat=VatOption.INCLUDE,
         )
         assert energy is not None
         assert isinstance(energy, Electricity)
@@ -70,6 +71,7 @@ async def test_electricity_none_date(aresponses: ResponsesMockServer) -> None:
         energy: Electricity = await client.energy_prices(
             start_date=today,
             end_date=today,
+            vat=VatOption.INCLUDE,
         )
         assert energy is not None
         assert isinstance(energy, Electricity)
@@ -96,6 +98,7 @@ async def test_electricity_midnight_cest(aresponses: ResponsesMockServer) -> Non
         energy: Electricity = await client.energy_prices(
             start_date=today,
             end_date=today,
+            vat=VatOption.INCLUDE,
         )
         assert energy is not None
         assert isinstance(energy, Electricity)
@@ -138,7 +141,11 @@ async def test_gas_model(aresponses: ResponsesMockServer) -> None:
     async with ClientSession() as session:
         today = date(2022, 12, 7)
         client = EnergyZero(session=session)
-        gas: Gas = await client.gas_prices(start_date=today, end_date=today)
+        gas: Gas = await client.gas_prices(
+            start_date=today,
+            end_date=today,
+            vat=VatOption.INCLUDE,
+        )
         assert gas is not None
         assert isinstance(gas, Gas)
         assert gas.extreme_prices[1] == 1.47
@@ -166,7 +173,9 @@ async def test_gas_morning_model(aresponses: ResponsesMockServer) -> None:
     async with ClientSession() as session:
         today = date(2022, 12, 7)
         client = EnergyZero(session=session)
-        gas: Gas = await client.gas_prices(start_date=today, end_date=today)
+        gas: Gas = await client.gas_prices(
+            start_date=today, end_date=today, vat=VatOption.INCLUDE
+        )
         assert gas is not None
         assert isinstance(gas, Gas)
         assert gas.current_price == 1.45