device class option + fixes

pilotak · Mar 6, 2021 · 07859bc · 07859bc
2 parents 50fdb90 + a32a831
commit 07859bc
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 28 deletions.
diff --git a/README.md b/README.md
@@ -12,7 +12,7 @@ sensor:
   - platform: attributes
     friendly_name: "Batteries"
     attribute: battery_level
-    unit_of_measurement: "%"
+    device_class: battery
     entities:
       - sensor.myslipo_1_0
       - sensor.myslipo_2_0
@@ -34,6 +34,7 @@ Configuration variables:
 - **attribute** (*Required*): Which attribute to extract from defined entity IDs.
 - **friendly_name** (*Optional*): Name to use in the Frontend *(will be the same for all entities specified)*.
 - **icon** (*Optional*): Icon to use in the Frontend.
+- **device_class** (*Optional*): Defines the device_class, if not specified it will be the same as parent.
 - **unit_of_measurement** (*Optional*): Defines the units of measurement of the sensor, if any.
 - **round_to** (*Optional*): Round numbers to 'x' decimals, if zero it will become whole number. Skip this field if you extracting a string or you want to leave the value as it is.
 - **value_template** (*Optional*): In case you need to do a math with the value ie. offset, bit gain, etc. *(will be the same for all entities specified)*.
@@ -71,7 +72,7 @@ sensor:
       - sensor.test3
 ```
 
->If an attribute is __`battery`__ or __`battery_level`__ and you don't specify __`icon`__, the following icon_template is applied (fullness). The result is that the battery icon becomes as full as the battery based on percentage.
+>If an attribute is __`battery`__ or __`battery_level`__ and you don't specify __`icon`__ or __`device_class`__ is not `battery`, the following icon_template is applied (fullness). The result is that the battery icon becomes as full as the battery based on percentage.
 
 ```yaml
 {% if batt == 'unknown' %}

diff --git a/custom_components/attributes/manifest.json b/custom_components/attributes/manifest.json
@@ -1,10 +1,10 @@
 {
   "domain": "attributes",
   "name": "Attributes",
+  "version": "1.1.1",
   "documentation": "https://github.com/pilotak/homeassistant-attributes",
+  "issue_tracker": "https://github.com/pilotak/homeassistant-attributes/issues",
   "requirements": [],
   "dependencies": [],
-  "codeowners": [
-  	"@pilotak"
-  ]
-}
+  "codeowners": ["@pilotak"]
+}
diff --git a/custom_components/attributes/sensor.py b/custom_components/attributes/sensor.py
@@ -10,7 +10,8 @@
 from homeassistant.components.sensor import ENTITY_ID_FORMAT, PLATFORM_SCHEMA
 from homeassistant.const import (
     ATTR_FRIENDLY_NAME, ATTR_UNIT_OF_MEASUREMENT, ATTR_ICON, CONF_ENTITIES,
-    EVENT_HOMEASSISTANT_START, STATE_UNKNOWN, CONF_VALUE_TEMPLATE)
+    ATTR_DEVICE_CLASS, EVENT_HOMEASSISTANT_START, STATE_UNKNOWN, STATE_UNAVAILABLE,
+    CONF_VALUE_TEMPLATE)
 from homeassistant.exceptions import TemplateError
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import Entity, async_generate_entity_id
@@ -30,6 +31,7 @@
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
     vol.Optional(ATTR_ICON): cv.string,
     vol.Optional(ATTR_FRIENDLY_NAME): cv.string,
+    vol.Optional(ATTR_DEVICE_CLASS): cv.string,
     vol.Optional(ATTR_UNIT_OF_MEASUREMENT): cv.string,
     vol.Optional(CONF_TIME_FORMAT): cv.string,
     vol.Optional(CONF_ROUND_TO): cv.positive_int,
@@ -52,28 +54,28 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
         if (attr == "last_triggered" or
                 attr == "last_changed") and time_format:
 
-            state_template = ("{{% if states('{0}') != '{2}' %}}\
-                              {{{{ as_timestamp(states.{0}.attributes['{1}'])\
+            state_template = ("{{% if states('{0}') != '{2}' and states('{0}') != '{4}' %}}\
+                              {{{{ as_timestamp(state_attr('{0}', '{1}'))\
                               | int | timestamp_local()\
                               | timestamp_custom('{2}') }}}}\
                               {{% else %}} {3} {{% endif %}}").format(
-                device, attr, time_format, STATE_UNKNOWN)
+                device, attr, time_format, STATE_UNKNOWN, STATE_UNAVAILABLE)
         else:
             round_to = config.get(CONF_ROUND_TO, None)
             additional_template = config.get(CONF_VALUE_TEMPLATE, "")
 
-            state_template = "{{% if states('{0}') != '{2}' %}}"
+            state_template = "{{% if states('{0}') != '{2}' and states('{0}') != '{5}' %}}"
 
             if round_to is None:
-                state_template += "{{{{ states.{0}.attributes['{1}'] {4} }}}}"
+                state_template += "{{{{ state_attr('{0}', '{1}') {4} }}}}"
             elif round_to > 0:
-                state_template += "{{{{ (states.{0}.attributes['{1}'] | float) | round({3}) {4} }}}}"
+                state_template += "{{{{ (state_attr('{0}', '{1}') | float) | round({3}) {4} }}}}"
             else:
-                state_template += "{{{{ states.{0}.attributes['{1}'] | int {4} }}}}"
+                state_template += "{{{{ state_attr('{0}', '{1}') | int {4} }}}}"
 
             state_template += "{{% else %}} {2} {{% endif %}}"
             state_template = state_template.format(
-                device, attr, STATE_UNKNOWN, round_to, additional_template)
+                device, attr, STATE_UNKNOWN, round_to, additional_template, STATE_UNAVAILABLE)
 
         _LOGGER.info("Adding attribute: %s of entity: %s", attr, device)
         _LOGGER.debug("Applying template: %s", state_template)
@@ -84,25 +86,33 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
         icon = str(config.get(ATTR_ICON))
 
         device_state = hass.states.get(device)
+
         if device_state is not None:
             device_friendly_name = device_state.attributes.get('friendly_name')
         else:
             device_friendly_name = device.split(".", 1)[1]
 
-        friendly_name = config.get(ATTR_FRIENDLY_NAME, device_friendly_name)
+        friendly_name = config.get(ATTR_FRIENDLY_NAME, None)
+
+        if device_state is not None:
+            device_class = config.get(
+                ATTR_DEVICE_CLASS, device_state.attributes.get('device_class'))
+        else:
+            device_class = config.get(ATTR_DEVICE_CLASS, None)
+
         unit_of_measurement = config.get(ATTR_UNIT_OF_MEASUREMENT)
 
-        if icon.startswith('mdi:'):
+        if icon.startswith('mdi:') or icon.startswith('hass:'):
             _LOGGER.debug("Applying user defined icon: '%s'", icon)
-            new_icon = ("{{% if states('{0}') != '{2}' %}} {1} {{% else %}}\
-                mdi:eye {{% endif %}}").format(device, icon, STATE_UNKNOWN)
+            new_icon = ("{{% if states('{0}') != '{2}' and states('{0}') != '{3}' %}} {1} {{% else %}}\
+                mdi:eye {{% endif %}}").format(device, icon, STATE_UNKNOWN, STATE_UNAVAILABLE)
 
             new_icon = template_helper.Template(new_icon)
             new_icon.hass = hass
-        elif attr == "battery" or attr == "battery_level":
+        elif (device_class is None or device_class != "battery") and attr == "battery" or attr == "battery_level":
             _LOGGER.debug("Applying battery icon template")
 
-            new_icon = ("{{% if states('{0}') != '{2}' %}}\
+            new_icon = ("{{% if states('{0}') != '{2}' and states('{0}') != '{3}' %}}\
                 {{% set batt = states.{0}.attributes['{1}']|int %}}\
                 {{% if batt == 'unknown' %}}\
                 mdi:battery-unknown\
@@ -131,7 +141,7 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
                 {{% endif %}}\
             {{% else %}}\
             mdi:battery-unknown\
-            {{% endif %}}").format(device, attr, STATE_UNKNOWN)
+            {{% endif %}}").format(device, attr, STATE_UNKNOWN, STATE_UNAVAILABLE)
             new_icon = template_helper.Template(str(new_icon))
             new_icon.hass = hass
         else:
@@ -143,6 +153,8 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
                 hass,
                 ("{0}_{1}").format(device.split(".", 1)[1], attr),
                 friendly_name,
+                device_friendly_name,
+                device_class,
                 unit_of_measurement,
                 state_template,
                 new_icon,
@@ -159,14 +171,16 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
 class AttributeSensor(RestoreEntity):
     """Representation of a Attribute Sensor."""
 
-    def __init__(self, hass, device_id, friendly_name, unit_of_measurement,
-                 state_template, icon_template, entity_id):
+    def __init__(self, hass, device_id, friendly_name, device_friendly_name, device_class,
+                 unit_of_measurement, state_template, icon_template, entity_id):
         """Initialize the sensor."""
         self.hass = hass
         self.entity_id = async_generate_entity_id(ENTITY_ID_FORMAT, device_id,
                                                   hass=hass)
-        self._name = friendly_name
+        self._name = friendly_name if friendly_name != None else device_friendly_name
+        self._friendly_name = friendly_name
         self._unique_id = slugify(f"{entity_id}_{device_id}")
+        self._device_class = device_class
         self._unit_of_measurement = unit_of_measurement
         self._template = state_template
         self._state = None
@@ -217,6 +231,11 @@ def icon(self):
         """Return the icon to use in the frontend, if any."""
         return self._icon
 
+    @property
+    def device_class(self):
+        """Return the device_class."""
+        return self._device_class
+
     @property
     def unit_of_measurement(self):
         """Return the unit_of_measurement of the device."""
@@ -232,7 +251,7 @@ def async_update(self):
         """Update the state from the template and the friendly name."""
 
         entity_state = self.hass.states.get(self._entity)
-        if entity_state is not None:
+        if self._friendly_name == None and entity_state != None:
             device_friendly_name = entity_state.attributes.get('friendly_name')
             if device_friendly_name is not None:
                 self._name = device_friendly_name

diff --git a/info.md b/info.md
@@ -9,7 +9,7 @@ sensor:
   - platform: attributes
     friendly_name: "Batteries"
     attribute: battery_level
-    unit_of_measurement: "%"
+    device_class: battery
     entities:
       - sensor.myslipo_1_0
       - sensor.myslipo_2_0
@@ -31,6 +31,7 @@ Configuration variables:
 - **attribute** (*Required*): Which attribute to extract from defined entity IDs.
 - **friendly_name** (*Optional*): Name to use in the Frontend *(will be the same for all entities specified)*.
 - **icon** (*Optional*): Icon to use in the Frontend.
+- **device_class** (*Optional*): Defines the device_class, if not specified it will be the same as parent.
 - **unit_of_measurement** (*Optional*): Defines the units of measurement of the sensor, if any.
 - **round_to** (*Optional*): Round numbers to 'x' decimals, if zero it will become whole number. Skip this field if you extracting a string or you want to leave the value as it is.
 - **value_template** (*Optional*): In case you need to do a math with the value ie. offset, bit gain, etc. *(will be the same for all entities specified)*.
@@ -50,7 +51,7 @@ sensor:
       - sensor.test3
 ```
 
->If an attribute is __`battery`__ or __`battery_level`__ and you don't specify __`icon`__, the following icon_template is applied (fullness). The result is that the battery icon becomes as full as the battery based on percentage.
+>If an attribute is __`battery`__ or __`battery_level`__ and you don't specify __`icon`__ or __`device_class`__ is not `battery`, the following icon_template is applied (fullness). The result is that the battery icon becomes as full as the battery based on percentage.
 
 ```yaml
 {% if batt == 'unknown' %}