tagged unions for events, commands and responses

Author: dmulcahey

tests/test_model.py

@@ -31,12 +31,14 @@ def test_ser_deser_zha_event():
         "device_ieee": "00:00:00:00:00:00:00:00",
         "unique_id": "00:00:00:00:00:00:00:00",
         "data": {"key": "value"},
+        "model_class_name": "ZHAEvent",
     }
 
     assert (
         zha_event.model_dump_json()
         == '{"message_type":"event","event_type":"device_event","event":"zha_event",'
-        '"device_ieee":"00:00:00:00:00:00:00:00","unique_id":"00:00:00:00:00:00:00:00","data":{"key":"value"}}'
+        '"device_ieee":"00:00:00:00:00:00:00:00","unique_id":"00:00:00:00:00:00:00:00",'
+        '"data":{"key":"value"},"model_class_name":"ZHAEvent"}'
     )
 
     device_info = DeviceInfo(

zha/application/websocket_api.py

@@ -9,12 +9,15 @@
 from pydantic import Field
 from zigpy.types.named import EUI64
 
-from zha.websocket.const import DURATION, GROUPS, APICommands
+from zha.websocket.const import GROUPS, APICommands
 from zha.websocket.server.api import decorators, register_api_command
 from zha.websocket.server.api.model import (
     GetApplicationStateResponse,
     GetDevicesResponse,
+    GroupsResponse,
+    PermitJoiningResponse,
     ReadClusterAttributesResponse,
+    UpdateGroupResponse,
     WebSocketCommand,
     WriteClusterAttributeResponse,
 )
@@ -150,7 +153,14 @@ async def get_groups(
             group.info_object
         )  # maybe we should change the group_id type...
     _LOGGER.info("groups: %s", groups)
-    client.send_result_success(command, {GROUPS: groups})
+    client.send_result_success(
+        command,
+        GroupsResponse(
+            **command.model_dump(exclude="model_class_name"),
+            groups=groups,
+            success=True,
+        ),
+    )
 
 
 class PermitJoiningCommand(WebSocketCommand):
@@ -169,10 +179,13 @@ async def permit_joining(
     """Permit joining devices to the Zigbee network."""
     # TODO add permit with code support
     await gateway.application_controller.permit(command.duration, command.ieee)
-    client.send_result_success(
-        command,
-        {DURATION: command.duration},
+    response = PermitJoiningResponse(
+        **command.model_dump(exclude="model_class_name"),
+        success=True,
+        duration=command.duration,
+        ieee=command.ieee,
     )
+    client.send_result_success(command, response)
 
 
 class RemoveDeviceCommand(WebSocketCommand):
@@ -358,7 +371,12 @@ async def create_group(
     members = command.members
     group_id = command.group_id
     group: Group = await gateway.async_create_zigpy_group(group_name, members, group_id)
-    client.send_result_success(command, {"group": group.info_object})
+    response = UpdateGroupResponse(
+        **command.model_dump(exclude="model_class_name"),
+        group=group.info_object,
+        success=True,
+    )
+    client.send_result_success(command, response)
 
 
 class RemoveGroupsCommand(WebSocketCommand):
@@ -416,7 +434,12 @@ async def add_group_members(
     if not group:
         client.send_result_error(command, "G1", "ZHA Group not found")
         return
-    client.send_result_success(command, {GROUP: group.info_object})
+    response = UpdateGroupResponse(
+        **command.model_dump(exclude="model_class_name"),
+        group=group.info_object,
+        success=True,
+    )
+    client.send_result_success(command, response)
 
 
 class RemoveGroupMembersCommand(AddGroupMembersCommand):
@@ -443,7 +466,12 @@ async def remove_group_members(
     if not group:
         client.send_result_error(command, "G1", "ZHA Group not found")
         return
-    client.send_result_success(command, {GROUP: group.info_object})
+    response = UpdateGroupResponse(
+        **command.model_dump(exclude="model_class_name"),
+        group=group.info_object,
+        success=True,
+    )
+    client.send_result_success(command, response)
 
 
 class StopServerCommand(WebSocketCommand):

zha/model.py

@@ -160,7 +160,7 @@ def as_tagged_union(union):
     ]
 
 
-class BaseEvent(BaseModel):
+class BaseEvent(TypedBaseModel):
     """Base model for ZHA events."""
 
     message_type: Literal["event"] = "event"

zha/websocket/client/model/messages.py

@@ -1,17 +1,11 @@
 """Models that represent messages in zha."""
 
-from typing import Annotated
-
 from pydantic import RootModel
-from pydantic.fields import Field
 
-from zha.websocket.server.api.model import CommandResponses, Events
+from zha.websocket.server.api.model import Messages
 
 
 class Message(RootModel):
     """Response model."""
 
-    root: Annotated[
-        CommandResponses | Events,
-        Field(discriminator="message_type"),
-    ]
+    root: Messages