Update throttle and add more attributes (#4644)

homeassistant/components/sensor/waqi.py

@@ -6,92 +6,129 @@ https://home-assistant.io/components/sensor.waqi/
 """
 import logging
 from datetime import timedelta
-from homeassistant.helpers.entity import Entity
-from homeassistant.util import Throttle
-from homeassistant.helpers.config_validation import PLATFORM_SCHEMA
-from homeassistant.helpers import config_validation as cv
+
 import voluptuous as vol
 
-REQUIREMENTS = ["pwaqi==1.3"]
+from homeassistant.const import (
+    ATTR_ATTRIBUTION, ATTR_TEMPERATURE, STATE_UNKNOWN)
+from homeassistant.helpers import config_validation as cv
+from homeassistant.helpers.config_validation import PLATFORM_SCHEMA
+from homeassistant.helpers.entity import Entity
+from homeassistant.util import Throttle
+
+REQUIREMENTS = ['pwaqi==1.3']
 
 _LOGGER = logging.getLogger(__name__)
 
+ATTR_DOMINENTPOL = 'dominentpol'
+ATTR_HUMIDITY = 'humidity'
+ATTR_NITROGEN_DIOXIDE = 'nitrogen_dioxide'
+ATTR_OZONE = 'ozone'
+ATTR_PARTICLE = 'particle'
+ATTR_PRESSURE = 'pressure'
+ATTR_TIME = 'time'
+ATTRIBUTION = 'Data provided by the World Air Quality Index project'
+
+CONF_LOCATIONS = 'locations'
+
+MIN_TIME_BETWEEN_UPDATES = timedelta(seconds=10)
+
 SENSOR_TYPES = {
     'aqi': ['AQI', '0-300+', 'mdi:cloud']
 }
 
-ATTR_LOCATION = 'locations'
-
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
-    vol.Required(ATTR_LOCATION): cv.ensure_list
+    vol.Required(CONF_LOCATIONS): cv.ensure_list
 })
 
-# Return cached results if last scan was less then this time ago
-MIN_TIME_BETWEEN_UPDATES = timedelta(minutes=30)
-
 
 def setup_platform(hass, config, add_devices, discovery_info=None):
-    """Setup the requested World Air Quality Index locations."""
-    dev = []
+    """Set up the requested World Air Quality Index locations."""
     import pwaqi
-    # Iterate each module
-    for location_name in config[ATTR_LOCATION]:
-        _LOGGER.debug('Adding location %s', location_name)
+
+    dev = []
+    for location_name in config.get(CONF_LOCATIONS):
         station_ids = pwaqi.findStationCodesByCity(location_name)
-        _LOGGER.debug('I got the following stations: %s', station_ids)
+        _LOGGER.error('The following stations were returned: %s', station_ids)
         for station in station_ids:
-            dev.append(WaqiSensor(station))
+            dev.append(WaqiSensor(WaqiData(station), station))
 
     add_devices(dev)
 
 
-# pylint: disable=too-few-public-methods
 class WaqiSensor(Entity):
     """Implementation of a WAQI sensor."""
 
-    def __init__(self, station_id):
+    def __init__(self, data, station_id):
         """Initialize the sensor."""
+        self.data = data
         self._station_id = station_id
-        self._state = None
+        self._details = None
         self.update()
 
     @property
     def name(self):
         """Return the name of the sensor."""
-        if 'city' in self._data:
-            return "WAQI {}".format(self._data['city']['name'])
-        return "WAQI {}".format(self._station_id)
+        try:
+            return 'WAQI {}'.format(self._details['city']['name'])
+        except (KeyError, TypeError):
+            return 'WAQI {}'.format(self._station_id)
 
     @property
     def icon(self):
         """Icon to use in the frontend, if any."""
-        return "mdi:cloud"
+        return 'mdi:cloud'
 
     @property
     def state(self):
         """Return the state of the device."""
-        return self._state
+        if self._details is not None:
+            return self._details.get('aqi')
+        else:
+            return STATE_UNKNOWN
 
     @property
     def unit_of_measurement(self):
         """Return the unit of measurement of this entity, if any."""
-        return "AQI"
+        return 'AQI'
 
     @property
     def state_attributes(self):
         """Return the state attributes of the last update."""
+        try:
             return {
-            "time": self._data.get('time', 'no data'),
-            "dominentpol": self._data.get('dominentpol', 'no data')
+                ATTR_ATTRIBUTION: ATTRIBUTION,
+                ATTR_TIME: self._details.get('time'),
+                ATTR_HUMIDITY: self._details['iaqi'][5]['cur'],
+                ATTR_PRESSURE: self._details['iaqi'][4]['cur'],
+                ATTR_TEMPERATURE: self._details['iaqi'][3]['cur'],
+                ATTR_OZONE: self._details['iaqi'][1]['cur'],
+                ATTR_PARTICLE: self._details['iaqi'][0]['cur'],
+                ATTR_NITROGEN_DIOXIDE: self._details['iaqi'][2]['cur'],
+                ATTR_DOMINENTPOL: self._details.get('dominentpol'),
             }
+        except (IndexError, KeyError):
+            return {ATTR_ATTRIBUTION: ATTRIBUTION}
+
+    def update(self):
+        """Get the latest data and updates the states."""
+        self.data.update()
+        self._details = self.data.data
+
+
+class WaqiData(object):
+    """Get the latest data and update the states."""
+
+    def __init__(self, station_id):
+        """Initialize the data object."""
+        self._station_id = station_id
+        self.data = None
 
     @Throttle(MIN_TIME_BETWEEN_UPDATES)
     def update(self):
         """Get the data from World Air Quality Index and updates the states."""
         import pwaqi
         try:
-            self._data = pwaqi.getStationObservation(self._station_id)
-
-            self._state = self._data.get('aqi', 'no data')
-        except KeyError:
-            _LOGGER.exception('Unable to fetch data from WAQI.')
+            self.data = pwaqi.getStationObservation(self._station_id)
+        except AttributeError:
+            _LOGGER.exception("Unable to fetch data from WAQI")