refactored to make its code structure similar to nsw_rural_fire_service_feed platform (#17461)

homeassistant/components/geo_location/geo_json_events.py

@@ -14,7 +14,10 @@ from homeassistant.components.geo_location import (
     PLATFORM_SCHEMA, GeoLocationEvent)
 from homeassistant.const import (
     CONF_RADIUS, CONF_SCAN_INTERVAL, CONF_URL, EVENT_HOMEASSISTANT_START)
+from homeassistant.core import callback
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.dispatcher import (
+    async_dispatcher_connect, dispatcher_send)
 from homeassistant.helpers.event import track_time_interval
 
 REQUIREMENTS = ['geojson_client==0.1']
@@ -28,6 +31,9 @@ DEFAULT_UNIT_OF_MEASUREMENT = 'km'
 
 SCAN_INTERVAL = timedelta(minutes=5)
 
+SIGNAL_DELETE_ENTITY = 'geo_json_events_delete_{}'
+SIGNAL_UPDATE_ENTITY = 'geo_json_events_update_{}'
+
 SOURCE = 'geo_json_events'
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
@@ -42,7 +48,14 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
     scan_interval = config.get(CONF_SCAN_INTERVAL, SCAN_INTERVAL)
     radius_in_km = config[CONF_RADIUS]
     # Initialize the entity manager.
-    GeoJsonFeedManager(hass, add_entities, scan_interval, url, radius_in_km)
+    feed = GeoJsonFeedManager(hass, add_entities, scan_interval, url,
+                              radius_in_km)
+
+    def start_feed_manager(event):
+        """Start feed manager."""
+        feed.startup()
+
+    hass.bus.listen_once(EVENT_HOMEASSISTANT_START, start_feed_manager)
 
 
 class GeoJsonFeedManager:
@@ -51,16 +64,19 @@ class GeoJsonFeedManager:
     def __init__(self, hass, add_entities, scan_interval, url, radius_in_km):
         """Initialize the GeoJSON Feed Manager."""
         from geojson_client.generic_feed import GenericFeed
+
         self._hass = hass
         self._feed = GenericFeed(
             (hass.config.latitude, hass.config.longitude),
             filter_radius=radius_in_km, url=url)
         self._add_entities = add_entities
         self._scan_interval = scan_interval
-        self._feed_entries = []
-        self._managed_entities = []
-        hass.bus.listen_once(
-            EVENT_HOMEASSISTANT_START, lambda _: self._update())
+        self.feed_entries = {}
+        self._managed_external_ids = set()
+
+    def startup(self):
+        """Start up this manager."""
+        self._update()
         self._init_regular_updates()
 
     def _init_regular_updates(self):
@@ -71,13 +87,24 @@ class GeoJsonFeedManager:
     def _update(self):
         """Update the feed and then update connected entities."""
         import geojson_client
+
         status, feed_entries = self._feed.update()
         if status == geojson_client.UPDATE_OK:
             _LOGGER.debug("Data retrieved %s", feed_entries)
             # Keep a copy of all feed entries for future lookups by entities.
-            self._feed_entries = feed_entries.copy()
-            keep_entries = self._update_or_remove_entities(feed_entries)
-            self._generate_new_entities(keep_entries)
+            self.feed_entries = {entry.external_id: entry
+                                 for entry in feed_entries}
+            # For entity management the external ids from the feed are used.
+            feed_external_ids = set(self.feed_entries)
+            remove_external_ids = self._managed_external_ids.difference(
+                feed_external_ids)
+            self._remove_entities(remove_external_ids)
+            update_external_ids = self._managed_external_ids.intersection(
+                feed_external_ids)
+            self._update_entities(update_external_ids)
+            create_external_ids = feed_external_ids.difference(
+                self._managed_external_ids)
+            self._generate_new_entities(create_external_ids)
         elif status == geojson_client.UPDATE_OK_NO_DATA:
             _LOGGER.debug(
                 "Update successful, but no data received from %s", self._feed)
@@ -85,61 +112,65 @@ class GeoJsonFeedManager:
             _LOGGER.warning(
                 "Update not successful, no data received from %s", self._feed)
             # Remove all entities.
-            self._update_or_remove_entities([])
+            self._remove_entities(self._managed_external_ids.copy())
 
-    def _update_or_remove_entities(self, feed_entries):
-        """Update existing entries and remove obsolete entities."""
-        _LOGGER.debug("Entries for updating: %s", feed_entries)
-        remove_entry = None
-        # Remove obsolete entities for events that have disappeared
-        managed_entities = self._managed_entities.copy()
-        for entity in managed_entities:
-            # Remove entry from previous iteration - if applicable.
-            if remove_entry:
-                feed_entries.remove(remove_entry)
-                remove_entry = None
-            for entry in feed_entries:
-                if entity.external_id == entry.external_id:
-                    # Existing entity - update details.
-                    _LOGGER.debug("Existing entity found %s", entity)
-                    remove_entry = entry
-                    entity.schedule_update_ha_state(True)
-                    break
-            else:
-                # Remove obsolete entity.
-                _LOGGER.debug("Entity not current anymore %s", entity)
-                self._managed_entities.remove(entity)
-                self._hass.add_job(entity.async_remove())
-        # Remove entry from very last iteration - if applicable.
-        if remove_entry:
-            feed_entries.remove(remove_entry)
-        # Return the remaining entries that new entities must be created for.
-        return feed_entries
-
-    def _generate_new_entities(self, entries):
+    def _generate_new_entities(self, external_ids):
         """Generate new entities for events."""
         new_entities = []
-        for entry in entries:
-            new_entity = GeoJsonLocationEvent(self, entry)
-            _LOGGER.debug("New entity added %s", new_entity)
+        for external_id in external_ids:
+            new_entity = GeoJsonLocationEvent(self, external_id)
+            _LOGGER.debug("New entity added %s", external_id)
             new_entities.append(new_entity)
-        # Add new entities to HA and keep track of them in this manager.
+            self._managed_external_ids.add(external_id)
+        # Add new entities to HA.
         self._add_entities(new_entities, True)
-        self._managed_entities.extend(new_entities)
 
-    def get_feed_entry(self, external_id):
-        """Return a feed entry identified by external id."""
-        return next((entry for entry in self._feed_entries
-                     if entry.external_id == external_id), None)
+    def _update_entities(self, external_ids):
+        """Update entities."""
+        for external_id in external_ids:
+            _LOGGER.debug("Existing entity found %s", external_id)
+            dispatcher_send(
+                self._hass, SIGNAL_UPDATE_ENTITY.format(external_id))
+
+    def _remove_entities(self, external_ids):
+        """Remove entities."""
+        for external_id in external_ids:
+            _LOGGER.debug("Entity not current anymore %s", external_id)
+            self._managed_external_ids.remove(external_id)
+            dispatcher_send(
+                self._hass, SIGNAL_DELETE_ENTITY.format(external_id))
 
 
 class GeoJsonLocationEvent(GeoLocationEvent):
     """This represents an external event with GeoJSON data."""
 
-    def __init__(self, feed_manager, feed_entry):
+    def __init__(self, feed_manager, external_id):
         """Initialize entity with data from feed entry."""
         self._feed_manager = feed_manager
-        self._update_from_feed(feed_entry)
+        self._external_id = external_id
+        self._name = None
+        self._distance = None
+        self._latitude = None
+        self._longitude = None
+
+    async def async_added_to_hass(self):
+        """Call when entity is added to hass."""
+        async_dispatcher_connect(
+            self.hass, SIGNAL_DELETE_ENTITY.format(self._external_id),
+            self._delete_callback)
+        async_dispatcher_connect(
+            self.hass, SIGNAL_UPDATE_ENTITY.format(self._external_id),
+            self._update_callback)
+
+    @callback
+    def _delete_callback(self):
+        """Remove this entity."""
+        self.hass.async_create_task(self.async_remove())
+
+    @callback
+    def _update_callback(self):
+        """Call update method."""
+        self.async_schedule_update_ha_state(True)
 
     @property
     def should_poll(self):
@@ -148,7 +179,8 @@ class GeoJsonLocationEvent(GeoLocationEvent):
 
     async def async_update(self):
         """Update this entity from the data held in the feed manager."""
-        feed_entry = self._feed_manager.get_feed_entry(self.external_id)
+        _LOGGER.debug("Updating %s", self._external_id)
+        feed_entry = self._feed_manager.feed_entries.get(self._external_id)
         if feed_entry:
             self._update_from_feed(feed_entry)
 
@@ -158,7 +190,6 @@ class GeoJsonLocationEvent(GeoLocationEvent):
         self._distance = feed_entry.distance_to_home
         self._latitude = feed_entry.coordinates[0]
         self._longitude = feed_entry.coordinates[1]
-        self.external_id = feed_entry.external_id
 
     @property
     def source(self) -> str:
@@ -194,6 +225,6 @@ class GeoJsonLocationEvent(GeoLocationEvent):
     def device_state_attributes(self):
         """Return the device state attributes."""
         attributes = {}
-        if self.external_id:
-            attributes[ATTR_EXTERNAL_ID] = self.external_id
+        if self._external_id:
+            attributes[ATTR_EXTERNAL_ID] = self._external_id
         return attributes