Add entity registry methods and models for WebSocket clients

New models for entity registry responses (EntityRegistryEntry,
EntityRegistryEntryExtended, EntityRegistryUpdateResult) and
list/get/update/remove methods on both sync and async WS clients.
Also adds configurable max_size param to WS client init (default 16MB)
to handle large responses like full entity registry lists.

Update CI Python version from 3.9 to 3.11.

Author: adamlogan73

.github/workflows/test-suite.yml

@@ -21,15 +21,15 @@ jobs:
       - name: Setup Python
         uses: actions/setup-python@v3
         with:
-          python-version: "3.9"
+          python-version: "3.11"
       - name: Checkout
         uses: actions/checkout@v3
         with:
           ref: ${{ github.event.pull_request.head.sha }}
       - name: Install uv
         uses: astral-sh/setup-uv@v4
       - name: Install Dependencies
-        run: uv sync --group styling
+        run: uv sync --group dev
       - name: Run Ruff format
         run: uv run ruff format homeassistant_api
       - name: Run Ruff linting

homeassistant_api/asyncwebsocket.py

@@ -25,6 +25,9 @@
 from homeassistant_api.models import State
 from homeassistant_api.models.config_entries import DisableEnableResult
 from homeassistant_api.models.config_entries import FlowResult
+from homeassistant_api.models.entity_registry import EntityRegistryEntry
+from homeassistant_api.models.entity_registry import EntityRegistryEntryExtended
+from homeassistant_api.models.entity_registry import EntityRegistryUpdateResult
 from homeassistant_api.models.states import Context
 from homeassistant_api.models.websocket import AuthInvalid
 from homeassistant_api.models.websocket import AuthOk
@@ -48,12 +51,12 @@
 class AsyncWebsocketClient(BaseWebsocketClient):
     _async_conn: ws.ClientConnection | None
 
-    def __init__(self, api_url: str, token: str) -> None:
-        super().__init__(api_url, token)
+    def __init__(self, api_url: str, token: str, *, max_size: int = 2**24) -> None:
+        super().__init__(api_url, token, max_size=max_size)
         self._async_conn = None
 
     async def __aenter__(self) -> Self:
-        self._async_conn = await ws.connect(self.api_url)
+        self._async_conn = await ws.connect(self.api_url, max_size=self.max_size)
         await self._async_conn.__aenter__()
         okay = await self.authentication_phase()
         logger.info("Authenticated with Home Assistant (%s)", okay.ha_version)
@@ -665,6 +668,64 @@ async def delete_entry_subentry(self, entry_id: str, subentry_id: str) -> None:
             ),
         )
 
+    # ── Entity Registry ─────────────────────────────────────────
+
+    async def list_entity_registry(self) -> tuple[EntityRegistryEntry, ...]:
+        """
+        List all entity registry entries.
+
+        Sends command :code:`{"type": "config/entity_registry/list", ...}`.
+        """
+        return tuple(
+            EntityRegistryEntry.from_json(entry)
+            for entry in await self.recv_result_list(
+                await self.send("config/entity_registry/list"),
+            )
+        )
+
+    async def get_entity_registry_entry(
+        self,
+        entity_id: str,
+    ) -> EntityRegistryEntryExtended:
+        """
+        Get a single entity registry entry.
+
+        Sends command :code:`{"type": "config/entity_registry/get", ...}`.
+        """
+        result = await self.recv_result_dict(
+            await self.send("config/entity_registry/get", entity_id=entity_id),
+        )
+        return EntityRegistryEntryExtended.from_json(result)
+
+    async def update_entity_registry_entry(
+        self,
+        entity_id: str,
+        **kwargs: Any,
+    ) -> EntityRegistryUpdateResult:
+        """
+        Update an entity registry entry.
+
+        Sends command :code:`{"type": "config/entity_registry/update", ...}`.
+        """
+        result = await self.recv_result_dict(
+            await self.send(
+                "config/entity_registry/update",
+                entity_id=entity_id,
+                **kwargs,
+            ),
+        )
+        return EntityRegistryUpdateResult.from_json(result)
+
+    async def remove_entity_registry_entry(self, entity_id: str) -> None:
+        """
+        Remove an entity from the entity registry.
+
+        Sends command :code:`{"type": "config/entity_registry/remove", ...}`.
+        """
+        await self.recv(
+            await self.send("config/entity_registry/remove", entity_id=entity_id),
+        )
+
     @contextlib.asynccontextmanager
     async def listen_config_entries(
         self,

homeassistant_api/basewebsocket.py

@@ -21,18 +21,20 @@ class BaseWebsocketClient:
 
     api_url: str
     token: str
+    max_size: int
     _id_counter: int
     _result_responses: dict[int, ResultResponse | None]
     _event_responses: dict[int, list[EventResponse]]
     _ping_responses: dict[int, PingResponse]
 
-    def __init__(self, api_url: str, token: str) -> None:
+    def __init__(self, api_url: str, token: str, *, max_size: int = 2**24) -> None:
         parsed = urlparse.urlparse(api_url)
         if parsed.scheme not in {"ws", "wss"}:
             msg = f"Unknown scheme {parsed.scheme} in {api_url}"
             raise ValueError(msg)
         self.api_url = api_url
         self.token = token.strip()
+        self.max_size = max_size
 
         self._id_counter = 0
         self._result_responses = {}  # id -> response

homeassistant_api/models/__init__.py

@@ -27,6 +27,12 @@
 from .entity import BaseGroup
 from .entity import Entity
 from .entity import Group
+from .entity_registry import EntityCategory
+from .entity_registry import EntityDisabledBy
+from .entity_registry import EntityHiddenBy
+from .entity_registry import EntityRegistryEntry
+from .entity_registry import EntityRegistryEntryExtended
+from .entity_registry import EntityRegistryUpdateResult
 from .events import AsyncEvent
 from .events import BaseEvent
 from .events import Event
@@ -57,6 +63,12 @@
     "DiscoveryKey",
     "Domain",
     "Entity",
+    "EntityCategory",
+    "EntityDisabledBy",
+    "EntityHiddenBy",
+    "EntityRegistryEntry",
+    "EntityRegistryEntryExtended",
+    "EntityRegistryUpdateResult",
     "Event",
     "FlowContext",
     "FlowResult",

homeassistant_api/models/entity_registry.py

@@ -0,0 +1,76 @@
+"""Models for Home Assistant entity registry responses."""
+
+from enum import Enum
+from typing import Any
+
+from pydantic import Field
+
+from .base import BaseModel
+from .base import DatetimeIsoField
+
+
+class EntityDisabledBy(str, Enum):
+    """What disabled an entity."""
+
+    CONFIG_ENTRY = "config_entry"
+    DEVICE = "device"
+    HASS = "hass"
+    INTEGRATION = "integration"
+    USER = "user"
+
+
+class EntityHiddenBy(str, Enum):
+    """What hid an entity."""
+
+    INTEGRATION = "integration"
+    USER = "user"
+
+
+class EntityCategory(str, Enum):
+    """Category of an entity."""
+
+    CONFIG = "config"
+    DIAGNOSTIC = "diagnostic"
+
+
+class EntityRegistryEntry(BaseModel):
+    """An entity registry entry as returned by ``config/entity_registry/list``."""
+
+    area_id: str | None = None
+    categories: dict[str, str] = Field(default_factory=dict)
+    config_entry_id: str | None = None
+    config_subentry_id: str | None = None
+    created_at: DatetimeIsoField
+    device_id: str | None = None
+    disabled_by: EntityDisabledBy | None = None
+    entity_category: EntityCategory | None = None
+    entity_id: str
+    has_entity_name: bool
+    hidden_by: EntityHiddenBy | None = None
+    icon: str | None = None
+    id: str
+    modified_at: DatetimeIsoField
+    name: str | None = None
+    options: dict[str, Any] = Field(default_factory=dict)
+    original_name: str | None = None
+    platform: str
+    translation_key: str | None = None
+    unique_id: str
+
+
+class EntityRegistryEntryExtended(EntityRegistryEntry):
+    """Extended entity registry entry as returned by ``config/entity_registry/get`` and ``update``."""
+
+    aliases: list[str] = Field(default_factory=list)
+    capabilities: dict[str, Any] | None = None
+    device_class: str | None = None
+    original_device_class: str | None = None
+    original_icon: str | None = None
+
+
+class EntityRegistryUpdateResult(BaseModel):
+    """Result from ``config/entity_registry/update``."""
+
+    entity_entry: EntityRegistryEntryExtended
+    reload_delay: int | None = None
+    require_restart: bool = False

homeassistant_api/websocket.py

@@ -25,6 +25,9 @@
 from homeassistant_api.models import State
 from homeassistant_api.models.config_entries import DisableEnableResult
 from homeassistant_api.models.config_entries import FlowResult
+from homeassistant_api.models.entity_registry import EntityRegistryEntry
+from homeassistant_api.models.entity_registry import EntityRegistryEntryExtended
+from homeassistant_api.models.entity_registry import EntityRegistryUpdateResult
 from homeassistant_api.models.states import Context
 from homeassistant_api.models.websocket import AuthInvalid
 from homeassistant_api.models.websocket import AuthOk
@@ -48,23 +51,15 @@
 class WebsocketClient(BaseWebsocketClient):
     _conn: ws.ClientConnection | None
 
-    def __init__(self, api_url: str, token: str) -> None:
-        super().__init__(api_url, token)
+    def __init__(self, api_url: str, token: str, *, max_size: int = 2**24) -> None:
+        super().__init__(api_url, token, max_size=max_size)
         self._conn = None
 
-        self._id_counter = 0
-        self._result_responses: dict[int, ResultResponse | None] = {}  # id -> response
-        self._event_responses: dict[
-            int,
-            list[EventResponse],
-        ] = {}  # id -> [response, ...]
-        self._ping_responses: dict[int, PingResponse] = {}  # id -> (sent, received)
-
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}({self.api_url!r})"
 
     def __enter__(self) -> Self:
-        self._conn = ws.connect(self.api_url)
+        self._conn = ws.connect(self.api_url, max_size=self.max_size)
         self._conn.__enter__()
         okay = self.authentication_phase()
         logger.info("Authenticated with Home Assistant (%s)", okay.ha_version)
@@ -649,6 +644,61 @@ def delete_entry_subentry(self, entry_id: str, subentry_id: str) -> None:
             ),
         )
 
+    # ── Entity Registry ─────────────────────────────────────────
+
+    def list_entity_registry(self) -> tuple[EntityRegistryEntry, ...]:
+        """
+        List all entity registry entries.
+
+        Sends command :code:`{"type": "config/entity_registry/list", ...}`.
+        """
+        return tuple(
+            EntityRegistryEntry.from_json(entry)
+            for entry in self.recv_result_list(
+                self.send("config/entity_registry/list"),
+            )
+        )
+
+    def get_entity_registry_entry(self, entity_id: str) -> EntityRegistryEntryExtended:
+        """
+        Get a single entity registry entry.
+
+        Sends command :code:`{"type": "config/entity_registry/get", ...}`.
+        """
+        result = self.recv_result_dict(
+            self.send("config/entity_registry/get", entity_id=entity_id),
+        )
+        return EntityRegistryEntryExtended.from_json(result)
+
+    def update_entity_registry_entry(
+        self,
+        entity_id: str,
+        **kwargs: Any,
+    ) -> EntityRegistryUpdateResult:
+        """
+        Update an entity registry entry.
+
+        Sends command :code:`{"type": "config/entity_registry/update", ...}`.
+        """
+        result = self.recv_result_dict(
+            self.send(
+                "config/entity_registry/update",
+                entity_id=entity_id,
+                **kwargs,
+            ),
+        )
+        return EntityRegistryUpdateResult.from_json(result)
+
+    def remove_entity_registry_entry(self, entity_id: str) -> None:
+        """
+        Remove an entity from the entity registry.
+
+        Sends command :code:`{"type": "config/entity_registry/remove", ...}`.
+        """
+        self.recv(
+            self.send("config/entity_registry/remove", entity_id=entity_id),
+        )
+
     @contextlib.contextmanager
     def listen_config_entries(
         self,