Implement websockets!

.vscode/settings.json

@@ -0,0 +1,7 @@
+{
+    "python.testing.pytestArgs": [
+        "tests"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
+}

README.md

@@ -16,7 +16,7 @@
 > However, it is not going to be deprecated according to [this comment](https://github.com/home-assistant/developers.home-assistant/pull/2150#pullrequestreview-2017433583)
 > But it is recommended to use the Websocket API for new integrations.
 
-Here is a quick example.
+### REST API Examples
 
 ```py
 from homeassistant_api import Client
@@ -25,14 +25,43 @@ with Client(
     '<API Server URL>', # i.e. 'http://homeassistant.local:8123/api/'
     '<Your Long Lived Access-Token>'
 ) as client:
-    light = client.trigger_service('light', 'turn_on', {'entity_id': 'light.living_room'})
+    light = client.trigger_service('light', 'turn_on', entity_id="light.living_room")
 ```
 
 All the methods also support async/await!
 Just prefix the method with `async_` and pass the `use_async=True` argument to the `Client` constructor.
 Then you can use the methods as coroutines
 (i.e. `await light.async_turn_on(...)`).
 
+```py
+import asyncio
+from homeassistant_api import Client
+
+async def main():
+    with Client(
+        '<REST API Server URL>', # i.e. 'http://homeassistant.local:8123/api/'
+        '<Your Long Lived Access-Token>',
+        use_async=True
+    ) as client:
+    light = await client.async_trigger_service('light', 'turn_on', entity_id="light.living_room")
+
+asyncio.run(main())
+```
+
+### Websocket API Example
+
+```py
+from homeassistant_api import WebsocketClient
+
+with WebsocketClient(
+    '<WS API Server URL>', # i.e. 'ws://homeassistant.local:8123/api/websocket'
+    '<Your Long Lived Access-Token>'
+) as ws_client:
+    light = ws_client.trigger_service('light', 'turn_on', entity_id="light.living_room")
+```
+
+> Note: The Websocket API is not yet supported in async/await mode.
+
 ## Documentation
 
 All documentation, API reference, contribution guidelines and pretty much everything else

compose.yml

@@ -1,5 +1,3 @@
-version: "3.9"
-
 services:
   server:
     image: "homeassistant/home-assistant:stable"
@@ -19,6 +17,7 @@ services:
       - server
     environment:
       HOMEASSISTANTAPI_URL: http://server:8123/api
+      HOMEASSISTANTAPI_WS_URL: ws://server:8123/api/websocket
       HOMEASSISTANTAPI_TOKEN: ${HOMEASSISTANTAPI_TOKEN}
       DELAY: 60
 

docs/api.rst

@@ -4,4 +4,4 @@ Code Reference
 .. automodule:: homeassistant_api
       :platform: Linux, Windows, MacOS
       :inherited-members:
-      :exclude-members: model_json_schema, model_copy, model_rebuild, model_dump, construct, copy, dict, from_orm, json, parse_file, parse_obj, parse_raw, parse_str, parse_url, schema, schema_json, schema_yaml, schema_yml, to_orm, update_forward_refs, validate, validate_file, validate_obj, validate_raw, validate_str, validate_url, model_validate_strings, model_validate_json, model_validate, model_post_init, model_parametrized_name, model_extra, model_fields_set, model_dump_json, model_construct, model_computed_fields
+      :exclude-members: model_json_schema, model_copy, model_rebuild, model_dump, construct, copy, dict, from_orm, json, parse_file, model_validate, parse_raw, parse_str, parse_url, schema, schema_json, schema_yaml, schema_yml, to_orm, update_forward_refs, validate, validate_file, validate_obj, validate_raw, validate_str, validate_url, model_validate_strings, model_validate_json, model_validate, model_post_init, model_parametrized_name, model_extra, model_fields_set, model_dump_json, model_construct, model_computed_fields

homeassistant_api/__init__.py

@@ -20,6 +20,7 @@
     "ParameterMissingError",
     "RequestError",
     "UnauthorizedError",
+    "WebsocketClient",
 )
 
 from .client import Client
@@ -36,6 +37,7 @@
 )
 from .models import Domain, Entity, Event, Group, History, LogbookEntry, Service, State
 from .processing import Processing
+from .websocket import WebsocketClient
 
 Domain.model_rebuild()
 Entity.model_rebuild()

homeassistant_api/client.py

@@ -1,9 +1,8 @@
 """Module containing the primary Client class."""
 
 import logging
-from typing import Any
 import urllib.parse as urlparse
-import warnings
+from typing import Any
 
 from .rawasyncclient import RawAsyncClient
 from .rawclient import RawClient
@@ -41,9 +40,5 @@ def __init__(
                 RawClient.__init__(
                     self, api_url, token, verify_ssl=verify_ssl, **kwargs
                 )
-            warnings.warn(
-                "The REST API is being phased out and will be removed in a far future release. Please use the WebSocket API instead.",
-                DeprecationWarning,
-            )
         else:
             raise ValueError(f"Unknown scheme {parsed.scheme} in {api_url}")

homeassistant_api/errors.py

@@ -1,6 +1,6 @@
 """Module for custom error classes"""
 
-from typing import Union
+from typing import Optional, Union
 
 
 class HomeassistantAPIError(Exception):
@@ -55,8 +55,8 @@ def __init__(self, status_code: int, content: Union[str, bytes]) -> None:
 class UnauthorizedError(HomeassistantAPIError):
     """Error raised when an invalid token in used to authenticate with homeassistant."""
 
-    def __init__(self) -> None:
-        super().__init__("Invalid authentication token")
+    def __init__(self, message: Optional[str] = None) -> None:
+        super().__init__(message or "Invalid authentication token")
 
 
 class EndpointNotFoundError(HomeassistantAPIError):

homeassistant_api/models/base.py

@@ -3,8 +3,8 @@
 from datetime import datetime
 from typing import Annotated
 
-from pydantic import ConfigDict, BaseModel as PydanticBaseModel, PlainSerializer
-
+from pydantic import BaseModel as PydanticBaseModel
+from pydantic import ConfigDict, PlainSerializer
 
 DatetimeIsoField = Annotated[
     datetime,

homeassistant_api/models/domains.py

@@ -12,19 +12,24 @@
 from .states import State
 
 if TYPE_CHECKING:
-    from homeassistant_api import Client
+    from homeassistant_api import Client, WebsocketClient
 
 
 class Domain(BaseModel):
     """Model representing the domain that services belong to."""
 
-    def __init__(self, *args, _client: Optional["Client"] = None, **kwargs) -> None:
+    def __init__(
+        self,
+        *args,
+        _client: Optional[Union["Client", "WebsocketClient"]] = None,
+        **kwargs,
+    ) -> None:
         super().__init__(*args, **kwargs)
         if _client is None:
             raise ValueError("No client passed.")
         object.__setattr__(self, "_client", _client)
 
-    _client: "Client"
+    _client: Union["Client", "WebsocketClient"]
     domain_id: str = Field(
         ...,
         description="The name of the domain that services belong to. "
@@ -36,7 +41,9 @@ def __init__(self, *args, _client: Optional["Client"] = None, **kwargs) -> None:
     )
 
     @classmethod
-    def from_json(cls, json: Dict[str, Any], client: "Client") -> "Domain":
+    def from_json(
+        cls, json: Dict[str, Any], client: Union["Client", "WebsocketClient"]
+    ) -> "Domain":
         """Constructs Domain and Service models from json data."""
         if "domain" not in json or "services" not in json:
             raise ValueError("Missing services or domain attribute in json argument.")
@@ -96,10 +103,15 @@ class Service(BaseModel):
     description: Optional[str] = None
     fields: Optional[Dict[str, ServiceField]] = None
 
-    def trigger(
-        self, **service_data
-    ) -> Union[Tuple[State, ...], Tuple[Tuple[State, ...], Dict[str, Any]]]:
+    def trigger(self, entity_id: Optional[str] = None, **service_data) -> Union[
+        Tuple[State, ...],
+        Tuple[Tuple[State, ...], Dict[str, Any]],
+        dict[str, Any],
+        None,
+    ]:
         """Triggers the service associated with this object."""
+        if entity_id is not None:
+            service_data["entity_id"] = entity_id
         try:
             return self.domain._client.trigger_service_with_response(
                 self.domain.domain_id,
@@ -114,9 +126,18 @@ def trigger(
             )
 
     async def async_trigger(
-        self, **service_data
+        self, entity_id: Optional[str] = None, **service_data
     ) -> Union[Tuple[State, ...], Tuple[Tuple[State, ...], Dict[str, Any]]]:
         """Triggers the service associated with this object."""
+        if entity_id is not None:
+            service_data["entity_id"] = entity_id
+
+        from homeassistant_api import WebsocketClient  # prevent circular import
+
+        if isinstance(self.domain._client, WebsocketClient):
+            raise NotImplementedError(
+                "WebsocketClient does not support async/await syntax."
+            )
         try:
             return await self.domain._client.async_trigger_service_with_response(
                 self.domain.domain_id,
@@ -130,8 +151,13 @@ async def async_trigger(
                 **service_data,
             )
 
-    def __call__(self, **service_data) -> Union[
-        Union[Tuple[State, ...], Tuple[Tuple[State, ...], Dict[str, Any]]],
+    def __call__(self, entity_id: Optional[str] = None, **service_data) -> Union[
+        Union[
+            Tuple[State, ...],
+            Tuple[Tuple[State, ...], Dict[str, Any]],
+            dict[str, Any],
+            None,
+        ],
         Coroutine[
             Any, Any, Union[Tuple[State, ...], Tuple[Tuple[State, ...], Dict[str, Any]]]
         ],
@@ -145,7 +171,7 @@ def __call__(self, **service_data) -> Union[
             if inspect.iscoroutinefunction(
                 caller := gc.get_referrers(parent_frame.f_code)[0]
             ) or inspect.iscoroutine(caller):
-                return self.async_trigger(**service_data)
+                return self.async_trigger(entity_id=entity_id, **service_data)
         except IndexError:  # pragma: no cover
             pass
-        return self.trigger(**service_data)
+        return self.trigger(entity_id=entity_id, **service_data)

homeassistant_api/models/states.py

@@ -12,9 +12,22 @@ class Context(BaseModel):
     """Model for entity state contexts."""
 
     id: str = Field(
-        max_length=128,
+        max_length=128,  # arbitrary limit
         description="Unique string identifying the context.",
     )
+    parent_id: Optional[str] = Field(
+        max_length=128,
+        description="Unique string identifying the parent context.",
+    )
+    user_id: Optional[str] = Field(
+        max_length=128,
+        description="Unique string identifying the user.",
+    )
+
+    @classmethod
+    def from_json(cls, json: Dict[str, Any]) -> "Context":
+        """Constructs Context model from json data"""
+        return cls.model_validate(json)
 
 
 class State(BaseModel):

homeassistant_api/models/websocket.py

@@ -0,0 +1,97 @@
+"""A module defining the responses we expect from the websocket API."""
+
+from typing import Any, Literal, Optional, Union
+
+from .base import BaseModel
+from .states import Context, DatetimeIsoField
+
+__all__ = (
+    "AuthRequired",
+    "AuthOk",
+    "AuthInvalid",
+    "PingResponse",
+    "ErrorResponse",
+    "ResultResponse",
+    "EventResponse",
+)
+
+
+class AuthRequired(BaseModel):
+    type: Literal["auth_required"]
+    ha_version: str
+
+
+class AuthOk(BaseModel):
+    type: Literal["auth_ok"]
+    ha_version: str
+
+
+class AuthInvalid(BaseModel):
+    type: Literal["auth_invalid"]
+    message: str
+
+
+class PingResponse(BaseModel):
+    """Ping websocket response model."""
+
+    id: int
+    type: Literal["pong"]
+    start: int  # added by the client, nanoseconds
+    end: Optional[int] = None  # added by the client, nanoseconds
+
+
+class Error(BaseModel):
+    code: str
+    message: str
+
+
+class ErrorResponse(BaseModel):
+    """Error websocket response model."""
+
+    id: int
+    success: Literal[False]
+    type: Literal["result"]
+    error: Error
+
+
+class ResultResponse(BaseModel):
+    """Result websocket response model."""
+
+    id: int
+    success: Literal[True]
+    type: Literal["result"]
+    result: Optional[Any]
+
+
+class FiredEvent(BaseModel):
+    """A model to parse the `event` key of fired event websocket responses."""
+
+    event_type: str
+    data: dict[str, Any]
+
+    origin: Literal["LOCAL", "REMOTE"]
+    # REMOTE if another API client or webhook fired the event
+    # LOCAL if Home Assistant (or the auth token we used) fired the event
+
+    time_fired: DatetimeIsoField  # datetime.datetime
+    context: Optional[Context]
+
+
+class TemplateEvent(BaseModel):
+    result: str
+    listeners: dict[str, Any]
+
+
+class FiredTrigger(BaseModel):
+    """A model to parse the `trigger` key of fired event websocket responses."""
+
+    context: Optional[Context]
+    variables: dict[str, Any]
+
+
+class EventResponse(BaseModel):
+    """A model to parse the response of a fired event websocket response."""
+
+    id: int
+    type: Literal["event"]
+    event: Union[FiredEvent, FiredTrigger, TemplateEvent]