WIP

Author: andrew-codechimp

custom_components/battery_notes/config_flow.py

@@ -10,7 +10,7 @@
 import homeassistant.helpers.entity_registry as er
 import voluptuous as vol
 from homeassistant import config_entries
-from homeassistant.components.sensor import SensorDeviceClass
+from homeassistant.components.sensor.const import SensorDeviceClass
 from homeassistant.config_entries import (
     ConfigEntry,
     ConfigFlowResult,
@@ -37,12 +37,11 @@
     CONF_MANUFACTURER,
     CONF_MODEL,
     CONF_MODEL_ID,
-    CONF_SHOW_ALL_DEVICES,
     CONF_SOURCE_ENTITY_ID,
     DATA_LIBRARY_UPDATER,
     DOMAIN,
-    DOMAIN_CONFIG,
 )
+from .coordinator import MY_KEY
 from .library import Library, ModelInfo
 from .library_updater import LibraryUpdater
 
@@ -216,9 +215,8 @@ async def async_step_device(
 
         schema = DEVICE_SCHEMA
         # If show_all_devices = is specified and true, don't filter
-        if DOMAIN in self.hass.data and DOMAIN_CONFIG in self.hass.data[DOMAIN]:
-            domain_config: dict = self.hass.data[DOMAIN][DOMAIN_CONFIG]
-            if domain_config.get(CONF_SHOW_ALL_DEVICES, False):
+        domain_config = self.hass.data[MY_KEY]
+        if domain_config.show_all_devices:
                 schema = DEVICE_SCHEMA_ALL
 
         return self.async_show_form(
@@ -359,6 +357,9 @@ async def async_step_battery(
             source_entity_id = self.data.get(CONF_SOURCE_ENTITY_ID, None)
             device_id = self.data.get(CONF_DEVICE_ID, None)
 
+            entity_entry = None
+            device_entry = None
+
             if source_entity_id:
                 entity_registry = er.async_get(self.hass)
                 entity_entry = entity_registry.async_get(source_entity_id)
@@ -372,6 +373,7 @@ async def async_step_battery(
                 unique_id = f"bn_{entity_unique_id}"
             else:
                 device_registry = dr.async_get(self.hass)
+                assert device_id
                 device_entry = device_registry.async_get(device_id)
                 unique_id = f"bn_{device_id}"
 
@@ -539,6 +541,7 @@ async def save_options(
 
         source_entity_id = self.config_entry.data.get(CONF_SOURCE_ENTITY_ID, None)
 
+        entity_entry: er.RegistryEntry | None = None
         if source_entity_id:
             entity_registry = er.async_get(self.hass)
             entity_entry = entity_registry.async_get(source_entity_id)

custom_components/battery_notes/coordinator.py

@@ -9,8 +9,8 @@
 
 from homeassistant.components.binary_sensor import DOMAIN as BINARY_SENSOR_DOMAIN
 from homeassistant.components.binary_sensor import BinarySensorDeviceClass
-from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
-from homeassistant.components.sensor import SensorDeviceClass
+from homeassistant.components.sensor.const import DOMAIN as SENSOR_DOMAIN
+from homeassistant.components.sensor.const import SensorDeviceClass
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_DEVICE_ID,
@@ -187,9 +187,8 @@ def __init__(
                 last_replaced,
             )
 
-            self.last_replaced = (
-                datetime.fromisoformat(last_replaced) if last_replaced else None
-            )
+            if last_replaced:
+                self.last_replaced = datetime.fromisoformat(last_replaced)
 
         # If there is not a last_reported set to now
         if not self.last_reported:
@@ -680,7 +679,7 @@ def battery_low(self) -> bool:
         if self.battery_low_template:
             return self.battery_low_template_state
         elif self.wrapped_battery:
-            if validate_is_float(self.current_battery_level):
+            if self.current_battery_level is not None and validate_is_float(self.current_battery_level):
                 return bool(
                     float(self.current_battery_level) < self.battery_low_threshold
                 )

custom_components/battery_notes/sensor.py

@@ -12,20 +12,26 @@
 from homeassistant.components.sensor import (
     PLATFORM_SCHEMA,
     RestoreSensor,
-    SensorDeviceClass,
     SensorEntity,
     SensorEntityDescription,
-    SensorStateClass,
 )
+from homeassistant.components.sensor.const import SensorDeviceClass, SensorStateClass
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_DEVICE_ID,
     CONF_NAME,
     PERCENTAGE,
     STATE_UNAVAILABLE,
     STATE_UNKNOWN,
+    EntityCategory,
+)
+from homeassistant.core import (
+    Event,
+    EventStateReportedData,
+    HomeAssistant,
+    callback,
+    split_entity_id,
 )
-from homeassistant.core import Event, HomeAssistant, callback, split_entity_id
 from homeassistant.helpers import (
     config_validation as cv,
 )
@@ -35,14 +41,13 @@
 from homeassistant.helpers import (
     entity_registry as er,
 )
-from homeassistant.helpers.entity import DeviceInfo, EntityCategory
+from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.entity_registry import (
     EVENT_ENTITY_REGISTRY_UPDATED,
 )
 from homeassistant.helpers.event import (
     EventStateChangedData,
-    EventStateReportedData,
     async_track_entity_registry_updated_event,
     async_track_state_change_event,
     async_track_state_report_event,
@@ -52,6 +57,7 @@
 from homeassistant.helpers.update_coordinator import (
     CoordinatorEntity,
 )
+from propcache.api import cached_property
 
 from .common import utcnow_no_timezone, validate_is_float
 from .const import (
@@ -241,7 +247,7 @@ async def async_setup_platform(
 
 
 class BatteryNotesBatteryPlusSensor(
-    RestoreSensor, SensorEntity, CoordinatorEntity[BatteryNotesCoordinator]
+    SensorEntity, CoordinatorEntity[BatteryNotesCoordinator]
 ):
     """Represents a battery plus type sensor."""
 
@@ -574,7 +580,7 @@ def _handle_coordinator_update(self) -> None:
 
         self.async_write_ha_state()
 
-    @property
+    @cached_property
     def extra_state_attributes(self) -> dict[str, Any] | None:
         """Return the state attributes of the battery type."""
 
@@ -604,7 +610,7 @@ def extra_state_attributes(self) -> dict[str, Any] | None:
             attrs.update(self._wrapped_attributes)
         return attrs
 
-    @property
+    @cached_property
     def native_value(self) -> StateType | Any | datetime:
         """Return the value reported by the sensor."""
         return self._attr_native_value