Improve zha websocket api logic (#68648)

This commit is contained in:
Marc Mueller 2022-03-25 17:45:47 +01:00 committed by GitHub
parent 931d5b5697
commit 6ac9c105c1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -334,17 +334,17 @@ async def websocket_get_device(
"""Get ZHA devices.""" """Get ZHA devices."""
zha_gateway: ZHAGateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY] zha_gateway: ZHAGateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
ieee: EUI64 = msg[ATTR_IEEE] ieee: EUI64 = msg[ATTR_IEEE]
device = None
if ieee in zha_gateway.devices: if not (zha_device := zha_gateway.devices.get(ieee)):
device = zha_gateway.devices[ieee].zha_device_info
if not device:
connection.send_message( connection.send_message(
websocket_api.error_message( websocket_api.error_message(
msg[ID], websocket_api.const.ERR_NOT_FOUND, "ZHA Device not found" msg[ID], websocket_api.const.ERR_NOT_FOUND, "ZHA Device not found"
) )
) )
return return
connection.send_result(msg[ID], device)
device_info = zha_device.zha_device_info
connection.send_result(msg[ID], device_info)
@websocket_api.require_admin @websocket_api.require_admin
@ -361,18 +361,17 @@ async def websocket_get_group(
"""Get ZHA group.""" """Get ZHA group."""
zha_gateway: ZHAGateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY] zha_gateway: ZHAGateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
group_id: int = msg[GROUP_ID] group_id: int = msg[GROUP_ID]
group = None
if group_id in zha_gateway.groups: if not (zha_group := zha_gateway.groups.get(group_id)):
group = zha_gateway.groups.get(group_id).group_info
if not group:
connection.send_message( connection.send_message(
websocket_api.error_message( websocket_api.error_message(
msg[ID], websocket_api.const.ERR_NOT_FOUND, "ZHA Group not found" msg[ID], websocket_api.const.ERR_NOT_FOUND, "ZHA Group not found"
) )
) )
return return
connection.send_result(msg[ID], group)
group_info = zha_group.group_info
connection.send_result(msg[ID], group_info)
def cv_group_member(value: Any) -> GroupMember: def cv_group_member(value: Any) -> GroupMember:
@ -453,18 +452,16 @@ async def websocket_add_group_members(
zha_gateway: ZHAGateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY] zha_gateway: ZHAGateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
group_id: int = msg[GROUP_ID] group_id: int = msg[GROUP_ID]
members: list[GroupMember] = msg[ATTR_MEMBERS] members: list[GroupMember] = msg[ATTR_MEMBERS]
zha_group = None
if group_id in zha_gateway.groups: if not (zha_group := zha_gateway.groups.get(group_id)):
zha_group = zha_gateway.groups.get(group_id)
await zha_group.async_add_members(members)
if not zha_group:
connection.send_message( connection.send_message(
websocket_api.error_message( websocket_api.error_message(
msg[ID], websocket_api.const.ERR_NOT_FOUND, "ZHA Group not found" msg[ID], websocket_api.const.ERR_NOT_FOUND, "ZHA Group not found"
) )
) )
return return
await zha_group.async_add_members(members)
ret_group = zha_group.group_info ret_group = zha_group.group_info
connection.send_result(msg[ID], ret_group) connection.send_result(msg[ID], ret_group)
@ -485,18 +482,16 @@ async def websocket_remove_group_members(
zha_gateway: ZHAGateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY] zha_gateway: ZHAGateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
group_id: int = msg[GROUP_ID] group_id: int = msg[GROUP_ID]
members: list[GroupMember] = msg[ATTR_MEMBERS] members: list[GroupMember] = msg[ATTR_MEMBERS]
zha_group = None
if group_id in zha_gateway.groups: if not (zha_group := zha_gateway.groups.get(group_id)):
zha_group = zha_gateway.groups.get(group_id)
await zha_group.async_remove_members(members)
if not zha_group:
connection.send_message( connection.send_message(
websocket_api.error_message( websocket_api.error_message(
msg[ID], websocket_api.const.ERR_NOT_FOUND, "ZHA Group not found" msg[ID], websocket_api.const.ERR_NOT_FOUND, "ZHA Group not found"
) )
) )
return return
await zha_group.async_remove_members(members)
ret_group = zha_group.group_info ret_group = zha_group.group_info
connection.send_result(msg[ID], ret_group) connection.send_result(msg[ID], ret_group)
@ -620,10 +615,8 @@ async def websocket_device_cluster_attributes(
endpoint_id, cluster_id, cluster_type endpoint_id, cluster_id, cluster_type
) )
if attributes is not None: if attributes is not None:
for attr_id in attributes: for attr_id, attr in attributes.items():
cluster_attributes.append( cluster_attributes.append({ID: attr_id, ATTR_NAME: attr[0]})
{ID: attr_id, ATTR_NAME: attributes[attr_id][0]}
)
_LOGGER.debug( _LOGGER.debug(
"Requested attributes for: %s: %s, %s: '%s', %s: %s, %s: %s", "Requested attributes for: %s: %s, %s: '%s', %s: %s, %s: %s",
ATTR_CLUSTER_ID, ATTR_CLUSTER_ID,
@ -660,7 +653,7 @@ async def websocket_device_cluster_commands(
cluster_id: int = msg[ATTR_CLUSTER_ID] cluster_id: int = msg[ATTR_CLUSTER_ID]
cluster_type: str = msg[ATTR_CLUSTER_TYPE] cluster_type: str = msg[ATTR_CLUSTER_TYPE]
zha_device = zha_gateway.get_device(ieee) zha_device = zha_gateway.get_device(ieee)
cluster_commands = [] cluster_commands: list[dict[str, Any]] = []
commands = None commands = None
if zha_device is not None: if zha_device is not None:
commands = zha_device.async_get_cluster_commands( commands = zha_device.async_get_cluster_commands(
@ -668,20 +661,20 @@ async def websocket_device_cluster_commands(
) )
if commands is not None: if commands is not None:
for cmd_id in commands[CLUSTER_COMMANDS_CLIENT]: for cmd_id, cmd in commands[CLUSTER_COMMANDS_CLIENT].items():
cluster_commands.append( cluster_commands.append(
{ {
TYPE: CLIENT, TYPE: CLIENT,
ID: cmd_id, ID: cmd_id,
ATTR_NAME: commands[CLUSTER_COMMANDS_CLIENT][cmd_id][0], ATTR_NAME: cmd[0],
} }
) )
for cmd_id in commands[CLUSTER_COMMANDS_SERVER]: for cmd_id, cmd in commands[CLUSTER_COMMANDS_SERVER].items():
cluster_commands.append( cluster_commands.append(
{ {
TYPE: CLUSTER_COMMAND_SERVER, TYPE: CLUSTER_COMMAND_SERVER,
ID: cmd_id, ID: cmd_id,
ATTR_NAME: commands[CLUSTER_COMMANDS_SERVER][cmd_id][0], ATTR_NAME: cmd[0],
} }
) )
_LOGGER.debug( _LOGGER.debug(