Update mvglive component (home-assistant#146479)

Co-authored-by: Erik Montnemery <[email protected]>

Author: danielpotthast

homeassistant/components/mvglive/manifest.json

@@ -2,10 +2,8 @@
   "domain": "mvglive",
   "name": "MVG",
   "codeowners": [],
-  "disabled": "This integration is disabled because it uses non-open source code to operate.",
   "documentation": "https://www.home-assistant.io/integrations/mvglive",
   "iot_class": "cloud_polling",
-  "loggers": ["MVGLive"],
-  "quality_scale": "legacy",
-  "requirements": ["PyMVGLive==1.1.4"]
+  "loggers": ["MVG"],
+  "requirements": ["mvg==1.4.0"]
 }

homeassistant/components/mvglive/sensor.py

@@ -1,13 +1,14 @@
 """Support for departure information for public transport in Munich."""
 
-# mypy: ignore-errors
 from __future__ import annotations
 
+from collections.abc import Mapping
 from copy import deepcopy
 from datetime import timedelta
 import logging
+from typing import Any
 
-import MVGLive
+from mvg import MvgApi, MvgApiError, TransportType
 import voluptuous as vol
 
 from homeassistant.components.sensor import (
@@ -19,6 +20,7 @@
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
+import homeassistant.util.dt as dt_util
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -44,53 +46,55 @@
     "SEV": "mdi:checkbox-blank-circle-outline",
     "-": "mdi:clock",
 }
-ATTRIBUTION = "Data provided by MVG-live.de"
+
+ATTRIBUTION = "Data provided by mvg.de"
 
 SCAN_INTERVAL = timedelta(seconds=30)
 
-PLATFORM_SCHEMA = SENSOR_PLATFORM_SCHEMA.extend(
-    {
-        vol.Required(CONF_NEXT_DEPARTURE): [
-            {
-                vol.Required(CONF_STATION): cv.string,
-                vol.Optional(CONF_DESTINATIONS, default=[""]): cv.ensure_list_csv,
-                vol.Optional(CONF_DIRECTIONS, default=[""]): cv.ensure_list_csv,
-                vol.Optional(CONF_LINES, default=[""]): cv.ensure_list_csv,
-                vol.Optional(
-                    CONF_PRODUCTS, default=DEFAULT_PRODUCT
-                ): cv.ensure_list_csv,
-                vol.Optional(CONF_TIMEOFFSET, default=0): cv.positive_int,
-                vol.Optional(CONF_NUMBER, default=1): cv.positive_int,
-                vol.Optional(CONF_NAME): cv.string,
-            }
-        ]
-    }
+PLATFORM_SCHEMA = vol.All(
+    cv.deprecated(CONF_DIRECTIONS),
+    SENSOR_PLATFORM_SCHEMA.extend(
+        {
+            vol.Required(CONF_NEXT_DEPARTURE): [
+                {
+                    vol.Required(CONF_STATION): cv.string,
+                    vol.Optional(CONF_DESTINATIONS, default=[""]): cv.ensure_list_csv,
+                    vol.Optional(CONF_DIRECTIONS, default=[""]): cv.ensure_list_csv,
+                    vol.Optional(CONF_LINES, default=[""]): cv.ensure_list_csv,
+                    vol.Optional(
+                        CONF_PRODUCTS, default=DEFAULT_PRODUCT
+                    ): cv.ensure_list_csv,
+                    vol.Optional(CONF_TIMEOFFSET, default=0): cv.positive_int,
+                    vol.Optional(CONF_NUMBER, default=1): cv.positive_int,
+                    vol.Optional(CONF_NAME): cv.string,
+                }
+            ]
+        }
+    ),
 )
 
 
-def setup_platform(
+async def async_setup_platform(
     hass: HomeAssistant,
     config: ConfigType,
     add_entities: AddEntitiesCallback,
     discovery_info: DiscoveryInfoType | None = None,
 ) -> None:
     """Set up the MVGLive sensor."""
-    add_entities(
-        (
-            MVGLiveSensor(
-                nextdeparture.get(CONF_STATION),
-                nextdeparture.get(CONF_DESTINATIONS),
-                nextdeparture.get(CONF_DIRECTIONS),
-                nextdeparture.get(CONF_LINES),
-                nextdeparture.get(CONF_PRODUCTS),
-                nextdeparture.get(CONF_TIMEOFFSET),
-                nextdeparture.get(CONF_NUMBER),
-                nextdeparture.get(CONF_NAME),
-            )
-            for nextdeparture in config[CONF_NEXT_DEPARTURE]
-        ),
-        True,
-    )
+    sensors = [
+        MVGLiveSensor(
+            hass,
+            nextdeparture.get(CONF_STATION),
+            nextdeparture.get(CONF_DESTINATIONS),
+            nextdeparture.get(CONF_LINES),
+            nextdeparture.get(CONF_PRODUCTS),
+            nextdeparture.get(CONF_TIMEOFFSET),
+            nextdeparture.get(CONF_NUMBER),
+            nextdeparture.get(CONF_NAME),
+        )
+        for nextdeparture in config[CONF_NEXT_DEPARTURE]
+    ]
+    add_entities(sensors, True)
 
 
 class MVGLiveSensor(SensorEntity):
@@ -100,38 +104,38 @@ class MVGLiveSensor(SensorEntity):
 
     def __init__(
         self,
-        station,
+        hass: HomeAssistant,
+        station_name,
         destinations,
-        directions,
         lines,
         products,
         timeoffset,
         number,
         name,
-    ):
+    ) -> None:
         """Initialize the sensor."""
-        self._station = station
         self._name = name
+        self._station_name = station_name
         self.data = MVGLiveData(
-            station, destinations, directions, lines, products, timeoffset, number
+            hass, station_name, destinations, lines, products, timeoffset, number
         )
         self._state = None
         self._icon = ICONS["-"]
 
     @property
-    def name(self):
+    def name(self) -> str | None:
         """Return the name of the sensor."""
         if self._name:
             return self._name
-        return self._station
+        return self._station_name
 
     @property
-    def native_value(self):
+    def native_value(self) -> str | None:
         """Return the next departure time."""
         return self._state
 
     @property
-    def extra_state_attributes(self):
+    def extra_state_attributes(self) -> Mapping[str, Any] | None:
         """Return the state attributes."""
         if not (dep := self.data.departures):
             return None
@@ -140,88 +144,114 @@ def extra_state_attributes(self):
         return attr
 
     @property
-    def icon(self):
+    def icon(self) -> str | None:
         """Icon to use in the frontend, if any."""
         return self._icon
 
     @property
-    def native_unit_of_measurement(self):
+    def native_unit_of_measurement(self) -> str | None:
         """Return the unit this state is expressed in."""
         return UnitOfTime.MINUTES
 
-    def update(self) -> None:
+    async def async_update(self) -> None:
         """Get the latest data and update the state."""
-        self.data.update()
+        await self.data.update()
         if not self.data.departures:
-            self._state = "-"
+            self._state = None
             self._icon = ICONS["-"]
         else:
-            self._state = self.data.departures[0].get("time", "-")
-            self._icon = ICONS[self.data.departures[0].get("product", "-")]
+            self._state = self.data.departures[0].get("time_in_mins", "-")
+            self._icon = self.data.departures[0].get("icon", ICONS["-"])
+
+
+def _get_minutes_until_departure(departure_time: int) -> int:
+    """Calculate the time difference in minutes between the current time and a given departure time.
+
+    Args:
+        departure_time: Unix timestamp of the departure time, in seconds.
+
+    Returns:
+        The time difference in minutes, as an integer.
+
+    """
+    current_time = dt_util.utcnow()
+    departure_datetime = dt_util.utc_from_timestamp(departure_time)
+    time_difference = (departure_datetime - current_time).total_seconds()
+    return int(time_difference / 60.0)
 
 
 class MVGLiveData:
-    """Pull data from the mvg-live.de web page."""
+    """Pull data from the mvg.de web page."""
 
     def __init__(
-        self, station, destinations, directions, lines, products, timeoffset, number
-    ):
+        self,
+        hass: HomeAssistant,
+        station_name,
+        destinations,
+        lines,
+        products,
+        timeoffset,
+        number,
+    ) -> None:
         """Initialize the sensor."""
-        self._station = station
+        self._hass = hass
+        self._station_name = station_name
+        self._station_id = None
         self._destinations = destinations
-        self._directions = directions
         self._lines = lines
         self._products = products
         self._timeoffset = timeoffset
         self._number = number
-        self._include_ubahn = "U-Bahn" in self._products
-        self._include_tram = "Tram" in self._products
-        self._include_bus = "Bus" in self._products
-        self._include_sbahn = "S-Bahn" in self._products
-        self.mvg = MVGLive.MVGLive()
-        self.departures = []
+        self.departures: list[dict[str, Any]] = []
 
-    def update(self):
+    async def update(self):
         """Update the connection data."""
+        if self._station_id is None:
+            try:
+                station = await MvgApi.station_async(self._station_name)
+                self._station_id = station["id"]
+            except MvgApiError as err:
+                _LOGGER.error(
+                    "Failed to resolve station %s: %s", self._station_name, err
+                )
+                self.departures = []
+                return
+
         try:
-            _departures = self.mvg.getlivedata(
-                station=self._station,
-                timeoffset=self._timeoffset,
-                ubahn=self._include_ubahn,
-                tram=self._include_tram,
-                bus=self._include_bus,
-                sbahn=self._include_sbahn,
+            _departures = await MvgApi.departures_async(
+                station_id=self._station_id,
+                offset=self._timeoffset,
+                limit=self._number,
+                transport_types=[
+                    transport_type
+                    for transport_type in TransportType
+                    if transport_type.value[0] in self._products
+                ]
+                if self._products
+                else None,
             )
         except ValueError:
             self.departures = []
             _LOGGER.warning("Returned data not understood")
             return
         self.departures = []
-        for i, _departure in enumerate(_departures):
-            # find the first departure meeting the criteria
+        for _departure in _departures:
             if (
                 "" not in self._destinations[:1]
                 and _departure["destination"] not in self._destinations
             ):
                 continue
 
-            if (
-                "" not in self._directions[:1]
-                and _departure["direction"] not in self._directions
-            ):
+            if "" not in self._lines[:1] and _departure["line"] not in self._lines:
                 continue
 
-            if "" not in self._lines[:1] and _departure["linename"] not in self._lines:
-                continue
+            time_to_departure = _get_minutes_until_departure(_departure["time"])
 
-            if _departure["time"] < self._timeoffset:
+            if time_to_departure < self._timeoffset:
                 continue
 
-            # now select the relevant data
             _nextdep = {}
-            for k in ("destination", "linename", "time", "direction", "product"):
+            for k in ("destination", "line", "type", "cancelled", "icon"):
                 _nextdep[k] = _departure.get(k, "")
-            _nextdep["time"] = int(_nextdep["time"])
+            _nextdep["time_in_mins"] = time_to_departure
             self.departures.append(_nextdep)
-            if i == self._number - 1:
-                break