Improve migration to device registry version 1.10 (home-assistant#151571)

Author: emontnemery

homeassistant/helpers/device_registry.py

@@ -465,7 +465,7 @@ class DeletedDeviceEntry:
         validator=_normalize_connections_validator
     )
     created_at: datetime = attr.ib()
-    disabled_by: DeviceEntryDisabler | None = attr.ib()
+    disabled_by: DeviceEntryDisabler | UndefinedType | None = attr.ib()
     id: str = attr.ib()
     identifiers: set[tuple[str, str]] = attr.ib()
     labels: set[str] = attr.ib()
@@ -480,15 +480,19 @@ def to_device_entry(
         config_subentry_id: str | None,
         connections: set[tuple[str, str]],
         identifiers: set[tuple[str, str]],
+        disabled_by: DeviceEntryDisabler | UndefinedType | None,
     ) -> DeviceEntry:
         """Create DeviceEntry from DeletedDeviceEntry."""
         # Adjust disabled_by based on config entry state
-        disabled_by = self.disabled_by
-        if config_entry.disabled_by:
-            if disabled_by is None:
-                disabled_by = DeviceEntryDisabler.CONFIG_ENTRY
-        elif disabled_by == DeviceEntryDisabler.CONFIG_ENTRY:
-            disabled_by = None
+        if self.disabled_by is not UNDEFINED:
+            disabled_by = self.disabled_by
+            if config_entry.disabled_by:
+                if disabled_by is None:
+                    disabled_by = DeviceEntryDisabler.CONFIG_ENTRY
+            elif disabled_by == DeviceEntryDisabler.CONFIG_ENTRY:
+                disabled_by = None
+        else:
+            disabled_by = disabled_by if disabled_by is not UNDEFINED else None
         return DeviceEntry(
             area_id=self.area_id,
             # type ignores: likely https://github.com/python/mypy/issues/8625
@@ -520,7 +524,10 @@ def as_storage_fragment(self) -> json_fragment:
                     },
                     "connections": list(self.connections),
                     "created_at": self.created_at,
-                    "disabled_by": self.disabled_by,
+                    "disabled_by": self.disabled_by
+                    if self.disabled_by is not UNDEFINED
+                    else None,
+                    "disabled_by_undefined": self.disabled_by is UNDEFINED,
                     "identifiers": list(self.identifiers),
                     "id": self.id,
                     "labels": list(self.labels),
@@ -621,6 +628,11 @@ async def _async_migrate_func(  # noqa: C901
                     device["connections"] = _normalize_connections(
                         device["connections"]
                     )
+            if old_minor_version < 12:
+                # Version 1.12 adds undefined flags to deleted devices, this is a bugfix
+                # of version 1.10
+                for device in old_data["deleted_devices"]:
+                    device["disabled_by_undefined"] = old_minor_version < 10
 
         if old_major_version > 2:
             raise NotImplementedError
@@ -934,6 +946,7 @@ def async_get_or_create(
                     config_subentry_id if config_subentry_id is not UNDEFINED else None,
                     connections,
                     identifiers,
+                    disabled_by,
                 )
                 disabled_by = UNDEFINED
 
@@ -1444,7 +1457,21 @@ async def async_load(self) -> None:
                     sw_version=device["sw_version"],
                     via_device_id=device["via_device_id"],
                 )
+
             # Introduced in 0.111
+            def get_optional_enum[_EnumT: StrEnum](
+                cls: type[_EnumT], value: str | None, undefined: bool
+            ) -> _EnumT | UndefinedType | None:
+                """Convert string to the passed enum, UNDEFINED or None."""
+                if undefined:
+                    return UNDEFINED
+                if value is None:
+                    return None
+                try:
+                    return cls(value)
+                except ValueError:
+                    return None
+
             for device in data["deleted_devices"]:
                 deleted_devices[device["id"]] = DeletedDeviceEntry(
                     area_id=device["area_id"],
@@ -1457,10 +1484,10 @@ async def async_load(self) -> None:
                     },
                     connections={tuple(conn) for conn in device["connections"]},
                     created_at=datetime.fromisoformat(device["created_at"]),
-                    disabled_by=(
-                        DeviceEntryDisabler(device["disabled_by"])
-                        if device["disabled_by"]
-                        else None
+                    disabled_by=get_optional_enum(
+                        DeviceEntryDisabler,
+                        device["disabled_by"],
+                        device["disabled_by_undefined"],
                     ),
                     identifiers={tuple(iden) for iden in device["identifiers"]},
                     id=device["id"],