Added days until pickup sensor (#9)

* Added days until pickup sensor

* Created a base class for both sensors

* Added types and comments

* Removed template sensor example

README.md

@@ -50,39 +50,3 @@ During the setup phase, you'll need to provide a URL from BIR.no
 3. Once you can see the dates, copy the URL from the address bar(should looks something similar to this: `https://bir.no/adressesoek/?rId=c2435f0f-2e4b-4908-86cf-bafbd3a2cf61&name=Lillehatten%20330,%20Bergen`).
 4. Paste this URL into the setup phase of the integration in Home Assistant's UI.
 
-## Tips: Changing Dates to 'Days Until Pickup'
-
-If you'd like to display the remaining days until the next pickup instead of the dates, you can create a [template sensor](https://www.home-assistant.io/integrations/template/) in Home Assistant.
-
-Here's an example:
-
-```yaml
-# configuration.yaml
-sensor:
-  - platform: template
-    sensors:
-      days_until_mixed_waste:
-        friendly_name: "Days Until Mixed Waste"
-        value_template: >-
-          {% set pick_up_date = as_timestamp(states('sensor.mixed_waste')) %}
-          {% set current_date = as_timestamp(now().strftime('%Y-%m-%d')) %}
-          {% set days_remaining = ((pick_up_date - current_date) / 86400) | int %}
-          {% if days_remaining >= 0 %}
-            {{ days_remaining }}
-          {% else %}
-            unknown
-          {% endif %}
-      days_until_paper_and_plastic_waste:
-        friendly_name: "Days Until Paper & Plastic Waste"
-        value_template: >-
-          {% set pick_up_date = as_timestamp(states('sensor.paper_and_plastic_waste')) %}
-          {% set current_date = as_timestamp(now().strftime('%Y-%m-%d')) %}
-          {% set days_remaining = ((pick_up_date - current_date) / 86400) | int %}
-          {% if days_remaining >= 0 %}
-            {{ days_remaining }}
-          {% else %}
-            unknown
-          {% endif %}
-
-
-```

custom_components/bir/sensor.py

@@ -1,76 +1,132 @@
-from datetime import timedelta, datetime
+from datetime import datetime, timedelta
 import aiohttp
 from homeassistant.components.sensor import SensorEntity
 import logging
-from .get_data import get_dates 
+from .get_data import get_dates
+from homeassistant.core import HomeAssistant
+from homeassistant.config_entries import ConfigEntry
+from typing import Any, Dict, List, Optional
+
 _LOGGER = logging.getLogger(__name__)
+
 SCAN_INTERVAL = timedelta(hours=1)
+NA_STRING = "N/A"
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
+async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities: Any) -> None:
+    """Set up the sensor platform.
+    
+    Args:
+        hass: HomeAssistant instance.
+        config_entry: Configuration entry for this sensor.
+        async_add_entities: Function to add entities to the platform.
+    """
     url = config_entry.data.get("url")
     session = aiohttp.ClientSession()
 
-    # Register a callback to close the session on Home Assistant's shutdown
-    async def close_session(event):
+    async def close_session(event: Any) -> None:
+        """Close the aiohttp session on Home Assistant stop event."""
         await session.close()
 
     hass.bus.async_listen_once("homeassistant_stop", close_session)
 
     data = await get_dates(session, url)
 
     if data:
-        sensors = []
+        sensors: List[SensorEntity] = []
         for waste_type, date in data.items():
-            sensor = WasteCollectionSensor(session, url, waste_type, date, config_entry.entry_id)
-            sensors.append(sensor)
+            collection_sensor = WasteCollectionSensorDates(session, url, waste_type, date, config_entry.entry_id)
+            days_until_sensor = WasteCollectionSensorDays(session, url, waste_type, date, config_entry.entry_id)
+            sensors.extend([collection_sensor, days_until_sensor])
 
-            # Update the sensor on the first run
-            await sensor.async_update()
+            await collection_sensor.async_update()
+            await days_until_sensor.async_update()
 
         if sensors:
-            async_add_entities(sensors, True)                
-
+            async_add_entities(sensors, True)
 
-class WasteCollectionSensor(SensorEntity):
-    def __init__(self, session, url, waste_type, state, entry_id):
+class WasteCollectionSensorBase(SensorEntity):
+    """Base sensor for waste collection."""
+
+    def __init__(self, session: aiohttp.ClientSession, url: str, waste_type: str, entry_id: str) -> None:
         self._session = session
         self._url = url
         self._waste_type = waste_type
-        self._state = state
         self._entry_id = entry_id
-        self._last_updated = None  # Initialize last updated attribute
+        self._last_updated: Optional[str] = None
 
     @property
-    def entity_registry_enabled_default(self) -> bool:
-        return True
+    def unique_id(self) -> str:
+        """Return the unique ID of the sensor."""
+        raise NotImplementedError
 
     @property
-    def unique_id(self):
-        return f"{self._entry_id}_{self._waste_type}"
+    def icon(self) -> str:
+        """Return the icon to be used for this sensor."""
+        return "mdi:trash-can"
 
     @property
-    def name(self):
-        return f"{self._waste_type.replace('_', ' ').title()}"
+    def extra_state_attributes(self) -> Dict[str, Any]:
+        """Return the state attributes of the sensor."""
+        return {"Last updated": self._last_updated}
+
+    async def async_update(self) -> None:
+        """Update the sensor state."""
+        data = await get_dates(self._session, self._url)
+        if data:
+            self._last_updated = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+class WasteCollectionSensorDates(WasteCollectionSensorBase):
+    """Sensor for showing waste collection dates."""
+
+    def __init__(self, session: aiohttp.ClientSession, url: str, waste_type: str, date: str, entry_id: str) -> None:
+        super().__init__(session, url, waste_type, entry_id)
+        self._date = date
+        self._state = NA_STRING
 
     @property
-    def state(self):
-        return self._state
-
+    def unique_id(self) -> str:
+        """Return the unique ID of the sensor."""
+        return f"{self._entry_id}_{self._waste_type}_date"
+
     @property
-    def icon(self):
-        """Return the icon of the sensor."""
-        # Example icon: mdi:recycle. Replace with your preferred icon
-        return "mdi:trash-can"
+    def name(self) -> str:
+        """Return the name of the sensor."""
+        return f"{self._waste_type.replace('_', ' ').title()} Collection Date"
 
     @property
-    def extra_state_attributes(self):
-        """Return the state attributes."""
-        return {
-            'Last updated': self._last_updated
-        }
+    def state(self) -> str:
+        """Return the state of the sensor."""
+        return self._state
 
-    async def async_update(self, *_):
+    async def async_update(self) -> None:
+        """Update the sensor state (pickup date)."""
         data = await get_dates(self._session, self._url)
-        if data:
-            self._state = data.get(self._waste_type, "N/A")
-            self._last_updated = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+        self._state = data.get(self._waste_type, NA_STRING) if data else NA_STRING
+
+class WasteCollectionSensorDays(WasteCollectionSensorBase):
+    """Sensor for showing days until the next waste collection."""
+
+    def __init__(self, session: aiohttp.ClientSession, url: str, waste_type: str, date: str, entry_id: str) -> None:
+        super().__init__(session, url, waste_type, entry_id)
+        self._date = date
+
+    @property
+    def unique_id(self) -> str:
+        """Return the unique ID of the sensor."""
+        return f"{self._entry_id}_{self._waste_type}_days"
+
+    @property
+    def name(self) -> str:
+        """Return the name of the sensor."""
+        return f"{self._waste_type.replace('_', ' ').title()} Days Until Pickup"
+
+    @property
+    def state(self) -> int:
+        """Return the state of the sensor."""
+        return self._calculate_days_until_pickup()
+
+    def _calculate_days_until_pickup(self) -> int:
+        """Calculate the days until the next waste collection."""
+        today = datetime.now().date()
+        pickup_date_obj = datetime.strptime(self._date, "%Y-%m-%d").date()
+        return (pickup_date_obj - today).days