Update sensor.py

finity69x2 · May 7, 2021 · bcd7dc1 · bcd7dc1
1 parent fc03fa8
commit bcd7dc1
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/custom_components/nws_alerts/sensor.py b/custom_components/nws_alerts/sensor.py
@@ -59,6 +59,7 @@ def __init__(self, name, zone_id):
         self._state = 0
         self._event = None
         self._event_id = None
+        self._message_type = None
         self._display_desc = None
         self._spoken_desc = None
         self._zone_id = zone_id.replace(' ', '')
@@ -93,6 +94,7 @@ def device_state_attributes(self):
         attrs[ATTR_ATTRIBUTION] = ATTRIBUTION
         attrs['title'] = self._event
         attrs['event_id'] = self._event_id
+        attrs['message_type'] = self._message_type
         attrs['display_desc'] = self._display_desc
         attrs['spoken_desc'] = self._spoken_desc
 
@@ -121,6 +123,7 @@ async def async_update(self):
         self._state = values['state']
         self._event = values['event']
         self._event_id = values['event_id']
+        self._message_type = values['message_type']
         self._display_desc = values['display_desc']
         self._spoken_desc = values['spoken_desc']
 
@@ -129,6 +132,7 @@ async def async_get_state(self):
             'state': self._state,
             'event': self._event,
             'event_id': self._event_id,
+            'message_type': self._message_type,
             'display_desc': self._display_desc,
             'spoken_desc': self._spoken_desc
         }
@@ -151,6 +155,7 @@ async def async_get_state(self):
                 'state': 0,
                 'event': None,
                 'event_id': None,
+                'message_type': None,
                 'display_desc': None,
                 'spoken_desc': None
             }
@@ -167,6 +172,7 @@ async def async_get_alerts(self):
             'state': self._state,
             'event': self._event,
             'event_id': self._event_id,
+            'message_type': self._message_type,
             'display_desc': self._display_desc,
             'spoken_desc': self._spoken_desc
         }
@@ -186,6 +192,7 @@ async def async_get_alerts(self):
             events = []
             headlines = []
             event_id = ''
+            message_type = ''
             display_desc = ''
             spoken_desc = ''
             features = data['features']
@@ -197,6 +204,7 @@ async def async_get_alerts(self):
                     headline = event
 
                 id = alert['id']
+                type = alert['properties']['messageType']
                 description = alert['properties']['description']
                 instruction = alert['properties']['instruction']
                 severity = alert['properties']['severity']
@@ -217,6 +225,8 @@ async def async_get_alerts(self):
                     event_id += '-'
 
                 event_id += id
+
+                message_type += type
 
             if headlines:
                 num_headlines = len(headlines)
@@ -241,15 +251,18 @@ async def async_get_alerts(self):
                 values['state'] = len(events)
                 values['event'] = event_str
                 values['event_id'] = event_id
+                values['message_type'] = message_type
                 values['display_desc'] = display_desc
                 values['spoken_desc'] = spoken_desc
             else:
                 values = {
                     'state': 0,
                     'event': None,
                     'event_id': None,
+                    'message_type': None,
                     'display_desc': None,
                     'spoken_desc': None
                 }
 
         return values
+