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

View File

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