Add HassKey compatibility with string keys

homeassistant/components/adguard/__init__.py

@@ -44,7 +44,7 @@
 )
 
 PLATFORMS = [Platform.SENSOR, Platform.SWITCH]
-ADGUARD_HASS_KEY = HassEntryKey["AdGuardData"]()
+ADGUARD_HASS_KEY = HassEntryKey["AdGuardData"](DOMAIN)
 
 
 @dataclass

homeassistant/setup.py

@@ -51,24 +51,26 @@
 #   being setup and the Task is the `_async_setup_component` helper.
 # - Tasks are removed from DATA_SETUP if setup was successful, that is,
 #   the task returned True.
-DATA_SETUP: HassKey[dict[str, asyncio.Future[bool]]] = HassKey()
+DATA_SETUP: HassKey[dict[str, asyncio.Future[bool]]] = HassKey("setup_tasks")
 
 # DATA_SETUP_DONE is a dict, indicating components which will be setup:
 # - Events are added to DATA_SETUP_DONE during bootstrap by
 #   async_set_domains_to_be_loaded, the key is the domain which will be loaded.
 # - Events are set and removed from DATA_SETUP_DONE when async_setup_component
 #   is finished, regardless of if the setup was successful or not.
-DATA_SETUP_DONE: HassKey[dict[str, asyncio.Future[bool]]] = HassKey()
+DATA_SETUP_DONE: HassKey[dict[str, asyncio.Future[bool]]] = HassKey("setup_done")
 
 # DATA_SETUP_STARTED is a dict, indicating when an attempt
 # to setup a component started.
-DATA_SETUP_STARTED: HassKey[dict[tuple[str, str | None], float]] = HassKey()
+DATA_SETUP_STARTED: HassKey[dict[tuple[str, str | None], float]] = HassKey(
+    "setup_started"
+)
 
 # DATA_SETUP_TIME is a defaultdict, indicating how time was spent
 # setting up a component.
 DATA_SETUP_TIME: HassKey[
     defaultdict[str, defaultdict[str | None, defaultdict[SetupPhases, float]]]
-] = HassKey()
+] = HassKey("setup_time")
 
 DATA_DEPS_REQS = "deps_reqs_processed"
 

homeassistant/util/hass_dict.py

@@ -2,20 +2,40 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Generic, TypeVar, assert_type, overload
 
 _T = TypeVar("_T")
 _U = TypeVar("_U")
 
 
+@dataclass(frozen=True)
 class _Key(Generic[_T]):
     """Base class for Hass key types."""
 
+    name: str
 
+    def __hash__(self) -> int:
+        """Return hash of name."""
+
+        return hash(self.name)
+
+    def __eq__(self, other: Any) -> bool:
+        """Check equality for dict keys to be compatible with str."""
+
+        if isinstance(other, str):
+            return self.name == other
+        if isinstance(other, _Key):
+            return self.name == other.name
+        return False
+
+
+@dataclass(frozen=True, eq=False)
 class HassEntryKey(_Key[_T]):
     """Key type for integrations with config entries."""
 
 
+@dataclass(frozen=True, eq=False)
 class HassKey(_Key[_T]):
     """Generic Hass key type."""
 
@@ -132,10 +152,10 @@ def test_hass_dict_typing() -> None:
         This is tested during the mypy run. Do not move it to 'tests'!
         """
         d = HassDict()
-        entry_key = HassEntryKey[int]()
-        key = HassKey[int]()
-        key2 = HassKey[dict[int, bool]]()
-        key3 = HassKey[set[str]]()
+        entry_key = HassEntryKey[int]("entry_key")
+        key = HassKey[int]("key")
+        key2 = HassKey[dict[int, bool]]("key2")
+        key3 = HassKey[set[str]]("key3")
         other_key = "domain"
 
         # __getitem__

tests/util/test_hass_dict.py

@@ -0,0 +1,47 @@
+"""Test HassDict and custom HassKey types."""
+
+from homeassistant.util.hass_dict import HassDict, HassEntryKey, HassKey
+
+
+def test_key_comparison() -> None:
+    """Test key comparison with itself and string keys."""
+
+    str_key = "custom-key"
+    key = HassKey[int](str_key)
+    other_key = HassKey[str]("other-key")
+
+    entry_key = HassEntryKey[int](str_key)
+    other_entry_key = HassEntryKey[str]("other-key")
+
+    assert key == str_key
+    assert key != other_key
+    assert key != 2
+
+    assert entry_key == str_key
+    assert entry_key != other_entry_key
+    assert entry_key != 2
+
+    # Only compare name attribute, HassKey(<name>) == HassEntryKey(<name>)
+    assert key == entry_key
+
+
+def test_hass_dict_access() -> None:
+    """Test keys with the same name all access the same value in HassDict."""
+
+    data = HassDict()
+    str_key = "custom-key"
+    key = HassKey[int](str_key)
+    other_key = HassKey[str]("other-key")
+
+    entry_key = HassEntryKey[int](str_key)
+    other_entry_key = HassEntryKey[str]("other-key")
+
+    data[str_key] = True
+    assert data.get(key) is True
+    assert data.get(other_key) is None
+
+    assert data.get(entry_key) is True  # type: ignore[comparison-overlap]
+    assert data.get(other_entry_key) is None
+
+    data[key] = False
+    assert data[str_key] is False