- Add different icons depending on transport type. Started with addin…

…g trams.
morosanmihail · Oct 18, 2023 · 37bf6a8 · 37bf6a8
1 parent 0d51a52
commit 37bf6a8
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/custom_components/london_tfl/const.py b/custom_components/london_tfl/const.py
@@ -1,8 +1,6 @@
 DOMAIN = 'london_tfl'
 
 DEFAULT_NAME = 'London TfL'
-DEFAULT_ICON = 'mdi:train'
-DEFAULT_BUS_ICON = 'mdi:bus'
 CONF_STOPS = 'stops'
 CONF_METHOD = 'method'
 CONF_LINE = 'line'
@@ -13,6 +11,7 @@
 DEFAULT_MAX = 3
 DEFAULT_LINES = {'dlr': 'DLR', 'jubilee': 'Jubilee'}
 DEFAULT_METHODS = ['tube', 'dlr', 'overground', 'cable-car', 'tram', 'river-tour', 'elizabeth-line', 'national-rail', 'bus']  # noqa
+DEFAULT_ICONS = {'bus': 'mdi:bus', 'tram': 'mdi:tram', 'default': 'mdi:train'}
 
 LINE_IMAGES = {
     'default': 'https://tfl.gov.uk/tfl/common/images/logos/London%20Underground/Roundel/LULRoundel.jpg',  # noqa

diff --git a/custom_components/london_tfl/sensor.py b/custom_components/london_tfl/sensor.py
@@ -16,9 +16,9 @@
 
 from .const import (
     CONF_STOPS, CONF_LINE, CONF_STATION, CONF_PLATFORM,
-    CONF_MAX, DEFAULT_ICON, DEFAULT_MAX, DEFAULT_NAME, DOMAIN,
+    CONF_MAX, DEFAULT_ICONS, DEFAULT_MAX, DEFAULT_NAME, DOMAIN,
     TFL_ARRIVALS_URL, CONF_SHORTEN_STATION_NAMES, get_line_image,
-    shortenName, CONF_METHOD, TFL_BUS_ARRIVALS_URL, DEFAULT_BUS_ICON
+    shortenName, CONF_METHOD, TFL_BUS_ARRIVALS_URL
 )
 from .network import request
 from .tfl_data import TfLData
@@ -150,7 +150,9 @@ def is_not_bus(self) -> bool:
     @property
     def icon(self):
         """Icon of the sensor."""
-        return DEFAULT_ICON if self.is_not_bus() else DEFAULT_BUS_ICON
+        if self.method in DEFAULT_ICONS:
+            return DEFAULT_ICONS[self.method]
+        return DEFAULT_ICONS['default']
 
     @property
     def state(self):