diff --git a/custom_components/london_tfl/const.py b/custom_components/london_tfl/const.py index f46def0..30ceec2 100644 --- 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 index 42befac..9117990 100644 --- 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):