Added city/state/country options and fixed several bugs

homeassistant/components/sensor/airvisual.py

@@ -13,9 +13,9 @@
 
 import homeassistant.helpers.config_validation as cv
 from homeassistant.components.sensor import PLATFORM_SCHEMA
-from homeassistant.const import (ATTR_ATTRIBUTION, ATTR_STATE, CONF_API_KEY,
-                                 CONF_LATITUDE, CONF_LONGITUDE,
-                                 CONF_MONITORED_CONDITIONS)
+from homeassistant.const import (ATTR_ATTRIBUTION, CONF_API_KEY, CONF_LATITUDE,
+                                 CONF_LONGITUDE, CONF_MONITORED_CONDITIONS,
+                                 CONF_STATE)
 from homeassistant.helpers.entity import Entity
 from homeassistant.util import Throttle
 
@@ -26,8 +26,11 @@
 ATTR_COUNTRY = 'country'
 ATTR_POLLUTANT_SYMBOL = 'pollutant_symbol'
 ATTR_POLLUTANT_UNIT = 'pollutant_unit'
+ATTR_REGION = 'region'
 ATTR_TIMESTAMP = 'timestamp'
 
+CONF_CITY = 'city'
+CONF_COUNTRY = 'country'
 CONF_RADIUS = 'radius'
 
 MASS_PARTS_PER_MILLION = 'ppm'
@@ -106,6 +109,12 @@
     cv.longitude,
     vol.Optional(CONF_RADIUS, default=1000):
     cv.positive_int,
+    vol.Optional(CONF_CITY):
+    cv.string,
+    vol.Optional(CONF_STATE):
+    cv.string,
+    vol.Optional(CONF_COUNTRY):
+    cv.string
 })
 
 
@@ -114,22 +123,28 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
     """Configure the platform and add the sensors."""
     import pyairvisual as pav
 
-    api_key = config.get(CONF_API_KEY)
-    _LOGGER.debug('AirVisual API Key: %s', api_key)
+    _LOGGER.debug('Received configuration: %s', config)
 
+    api_key = config.get(CONF_API_KEY)
     monitored_locales = config.get(CONF_MONITORED_CONDITIONS)
-    _LOGGER.debug('Monitored Conditions: %s', monitored_locales)
-
     latitude = config.get(CONF_LATITUDE, hass.config.latitude)
-    _LOGGER.debug('AirVisual Latitude: %s', latitude)
-
     longitude = config.get(CONF_LONGITUDE, hass.config.longitude)
-    _LOGGER.debug('AirVisual Longitude: %s', longitude)
-
     radius = config.get(CONF_RADIUS)
-    _LOGGER.debug('AirVisual Radius: %s', radius)
-
-    data = AirVisualData(pav.Client(api_key), latitude, longitude, radius)
+    city = config.get(CONF_CITY)
+    state = config.get(CONF_STATE)
+    country = config.get(CONF_COUNTRY)
+
+    if city and state and country:
+        _LOGGER.debug('Constructing sensors based on city, state, and country')
+        data = AirVisualData(
+            pav.Client(api_key), city=city, state=state, country=country)
+    else:
+        _LOGGER.debug('Constructing sensors based on latitude and longitude')
+        data = AirVisualData(
+            pav.Client(api_key),
+            latitude=latitude,
+            longitude=longitude,
+            radius=radius)
 
     sensors = []
     for locale in monitored_locales:
@@ -161,14 +176,13 @@ def __init__(self, data, name, icon, locale):
     @property
     def device_state_attributes(self):
         """Return the state attributes."""
-        if self._data:
-            return {
-                ATTR_ATTRIBUTION: 'AirVisual©',
-                ATTR_CITY: self._data.city,
-                ATTR_COUNTRY: self._data.country,
-                ATTR_STATE: self._data.state,
-                ATTR_TIMESTAMP: self._data.pollution_info.get('ts')
-            }
+        return {
+            ATTR_ATTRIBUTION: 'AirVisual©',
+            ATTR_CITY: self._data.city,
+            ATTR_COUNTRY: self._data.country,
+            ATTR_REGION: self._data.state,
+            ATTR_TIMESTAMP: self._data.pollution_info.get('ts')
+        }
 
     @property
     def icon(self):
@@ -188,7 +202,7 @@ def state(self):
     @asyncio.coroutine
     def async_update(self):
         """Update the status of the sensor."""
-        _LOGGER.debug('updating sensor: %s', self._name)
+        _LOGGER.debug('Updating sensor: %s', self._name)
         self._data.update()
 
 
@@ -200,13 +214,14 @@ def async_update(self):
         """Update the status of the sensor."""
         yield from super().async_update()
         aqi = self._data.pollution_info.get('aqi{0}'.format(self._locale))
-
         try:
             [level] = [
                 i for i in POLLUTANT_LEVEL_MAPPING
                 if i['minimum'] <= aqi <= i['maximum']
             ]
             self._state = level.get('label')
+        except TypeError:
+            self._state = None
         except ValueError:
             self._state = None
 
@@ -217,12 +232,13 @@ class AirQualityIndexSensor(AirVisualBaseSensor):
     @property
     def unit_of_measurement(self):
         """Return the unit the value is expressed in."""
-        return ''
+        return 'PSI'
 
     @asyncio.coroutine
     def async_update(self):
         """Update the status of the sensor."""
         yield from super().async_update()
+
         self._state = self._data.pollution_info.get(
             'aqi{0}'.format(self._locale))
 
@@ -239,11 +255,10 @@ def __init__(self, data, name, icon, locale):
     @property
     def device_state_attributes(self):
         """Return the state attributes."""
-        if self._data:
-            return merge_two_dicts(super().device_state_attributes, {
-                ATTR_POLLUTANT_SYMBOL: self._symbol,
-                ATTR_POLLUTANT_UNIT: self._unit
-            })
+        return merge_two_dicts(super().device_state_attributes, {
+            ATTR_POLLUTANT_SYMBOL: self._symbol,
+            ATTR_POLLUTANT_UNIT: self._unit
+        })
 
     @asyncio.coroutine
     def async_update(self):
@@ -259,31 +274,35 @@ def async_update(self):
 class AirVisualData(object):
     """Define an object to hold sensor data."""
 
-    def __init__(self, client, latitude, longitude, radius):
+    def __init__(self, client, **kwargs):
         """Initialize."""
-        self.city = None
         self._client = client
-        self.country = None
-        self.latitude = latitude
-        self.longitude = longitude
         self.pollution_info = None
-        self.radius = radius
-        self.state = None
+
+        self.city = kwargs.get(CONF_CITY)
+        self.state = kwargs.get(CONF_STATE)
+        self.country = kwargs.get(CONF_COUNTRY)
+
+        self.latitude = kwargs.get(CONF_LATITUDE)
+        self.longitude = kwargs.get(CONF_LONGITUDE)
+        self.radius = kwargs.get(CONF_RADIUS)
 
     @Throttle(MIN_TIME_BETWEEN_UPDATES)
     def update(self):
         """Update with new AirVisual data."""
         import pyairvisual.exceptions as exceptions
 
         try:
-            resp = self._client.nearest_city(self.latitude, self.longitude,
-                                             self.radius).get('data')
+            if self.city and self.state and self.country:
+                resp = self._client.city(self.city, self.state,
+                                         self.country).get('data')
+            else:
+                resp = self._client.nearest_city(self.latitude, self.longitude,
+                                                 self.radius).get('data')
             _LOGGER.debug('New data retrieved: %s', resp)
-
-            self.city = resp.get('city')
-            self.state = resp.get('state')
-            self.country = resp.get('country')
-            self.pollution_info = resp.get('current').get('pollution')
+            self.pollution_info = resp.get('current', {}).get('pollution', {})
         except exceptions.HTTPError as exc_info:
-            _LOGGER.error('Unable to update sensor data')
+            _LOGGER('Unable to retrieve data from the API')
+            _LOGGER.error("There is likely no data on this location")
             _LOGGER.debug(exc_info)
+            self.pollution_info = {}