Add websocket type hints in entity_registry (#80657)

* Add websocket type hints in entity_registry

* Adjust websocket_list_entities

* Fix update

* Fix websocket_update_entity

* Apply suggestion

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* Apply suggestion

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
This commit is contained in:
epenet 2022-10-20 12:20:39 +02:00 committed by GitHub
parent cce4485fb7
commit b23a66d776
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 29 additions and 10 deletions

View File

@ -36,14 +36,18 @@ async def async_setup(hass: HomeAssistant) -> bool:
{vol.Required("type"): "config/entity_registry/list"}
)
@callback
def websocket_list_entities(hass, connection, msg):
def websocket_list_entities(
hass: HomeAssistant,
connection: websocket_api.ActiveConnection,
msg: dict[str, Any],
) -> None:
"""Handle list registry entries command."""
nonlocal cached_list_entities
if not cached_list_entities:
registry = er.async_get(hass)
cached_list_entities = message_to_json(
websocket_api.result_message(
IDEN_TEMPLATE,
IDEN_TEMPLATE, # type: ignore[arg-type]
[_entry_dict(entry) for entry in registry.entities.values()],
)
)
@ -70,7 +74,11 @@ async def async_setup(hass: HomeAssistant) -> bool:
}
)
@callback
def websocket_get_entity(hass, connection, msg):
def websocket_get_entity(
hass: HomeAssistant,
connection: websocket_api.ActiveConnection,
msg: dict[str, Any],
) -> None:
"""Handle get entity registry entry command.
Async friendly.
@ -120,7 +128,11 @@ def websocket_get_entity(hass, connection, msg):
}
)
@callback
def websocket_update_entity(hass, connection, msg):
def websocket_update_entity(
hass: HomeAssistant,
connection: websocket_api.ActiveConnection,
msg: dict[str, Any],
) -> None:
"""Handle update entity websocket command.
Async friendly.
@ -153,7 +165,7 @@ def websocket_update_entity(hass, connection, msg):
if entity_entry.device_id:
device_registry = dr.async_get(hass)
device = device_registry.async_get(entity_entry.device_id)
if device.disabled:
if device and device.disabled:
connection.send_message(
websocket_api.error_message(
msg["id"], "invalid_info", "Device is disabled"
@ -184,11 +196,14 @@ def websocket_update_entity(hass, connection, msg):
)
return
result = {"entity_entry": _entry_ext_dict(entity_entry)}
result: dict[str, Any] = {"entity_entry": _entry_ext_dict(entity_entry)}
if "disabled_by" in changes and changes["disabled_by"] is None:
# Enabling an entity requires a config entry reload, or HA restart
config_entry = hass.config_entries.async_get_entry(entity_entry.config_entry_id)
if config_entry and not config_entry.supports_unload:
if (
not (config_entry_id := entity_entry.config_entry_id)
or (config_entry := hass.config_entries.async_get_entry(config_entry_id))
and not config_entry.supports_unload
):
result["require_restart"] = True
else:
result["reload_delay"] = config_entries.RELOAD_AFTER_UPDATE_DELAY
@ -203,7 +218,11 @@ def websocket_update_entity(hass, connection, msg):
}
)
@callback
def websocket_remove_entity(hass, connection, msg):
def websocket_remove_entity(
hass: HomeAssistant,
connection: websocket_api.ActiveConnection,
msg: dict[str, Any],
) -> None:
"""Handle remove entity websocket command.
Async friendly.

View File

@ -345,7 +345,7 @@ async def test_update_entity(hass, client):
"platform": "test_platform",
"unique_id": "1234",
},
"reload_delay": 30,
"require_restart": True,
}
# UPDATE ENTITY OPTION