Replace functools.partial with ServiceCall.hass in unifiprotect (#133131)

This commit is contained in:
epenet 2024-12-13 13:28:54 +01:00 committed by GitHub
parent f816a0667c
commit c7adc98408
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -3,7 +3,6 @@
from __future__ import annotations
import asyncio
import functools
from typing import Any, cast
from pydantic.v1 import ValidationError
@ -88,9 +87,9 @@ def _async_get_ufp_instance(hass: HomeAssistant, device_id: str) -> ProtectApiCl
@callback
def _async_get_ufp_camera(hass: HomeAssistant, call: ServiceCall) -> Camera:
ref = async_extract_referenced_entity_ids(hass, call)
entity_registry = er.async_get(hass)
def _async_get_ufp_camera(call: ServiceCall) -> Camera:
ref = async_extract_referenced_entity_ids(call.hass, call)
entity_registry = er.async_get(call.hass)
entity_id = ref.indirectly_referenced.pop()
camera_entity = entity_registry.async_get(entity_id)
@ -98,30 +97,27 @@ def _async_get_ufp_camera(hass: HomeAssistant, call: ServiceCall) -> Camera:
assert camera_entity.device_id is not None
camera_mac = _async_unique_id_to_mac(camera_entity.unique_id)
instance = _async_get_ufp_instance(hass, camera_entity.device_id)
instance = _async_get_ufp_instance(call.hass, camera_entity.device_id)
return cast(Camera, instance.bootstrap.get_device_from_mac(camera_mac))
@callback
def _async_get_protect_from_call(
hass: HomeAssistant, call: ServiceCall
) -> set[ProtectApiClient]:
def _async_get_protect_from_call(call: ServiceCall) -> set[ProtectApiClient]:
return {
_async_get_ufp_instance(hass, device_id)
_async_get_ufp_instance(call.hass, device_id)
for device_id in async_extract_referenced_entity_ids(
hass, call
call.hass, call
).referenced_devices
}
async def _async_service_call_nvr(
hass: HomeAssistant,
call: ServiceCall,
method: str,
*args: Any,
**kwargs: Any,
) -> None:
instances = _async_get_protect_from_call(hass, call)
instances = _async_get_protect_from_call(call)
try:
await asyncio.gather(
*(getattr(i.bootstrap.nvr, method)(*args, **kwargs) for i in instances)
@ -130,23 +126,23 @@ async def _async_service_call_nvr(
raise HomeAssistantError(str(err)) from err
async def add_doorbell_text(hass: HomeAssistant, call: ServiceCall) -> None:
async def add_doorbell_text(call: ServiceCall) -> None:
"""Add a custom doorbell text message."""
message: str = call.data[ATTR_MESSAGE]
await _async_service_call_nvr(hass, call, "add_custom_doorbell_message", message)
await _async_service_call_nvr(call, "add_custom_doorbell_message", message)
async def remove_doorbell_text(hass: HomeAssistant, call: ServiceCall) -> None:
async def remove_doorbell_text(call: ServiceCall) -> None:
"""Remove a custom doorbell text message."""
message: str = call.data[ATTR_MESSAGE]
await _async_service_call_nvr(hass, call, "remove_custom_doorbell_message", message)
await _async_service_call_nvr(call, "remove_custom_doorbell_message", message)
async def remove_privacy_zone(hass: HomeAssistant, call: ServiceCall) -> None:
async def remove_privacy_zone(call: ServiceCall) -> None:
"""Remove privacy zone from camera."""
name: str = call.data[ATTR_NAME]
camera = _async_get_ufp_camera(hass, call)
camera = _async_get_ufp_camera(call)
remove_index: int | None = None
for index, zone in enumerate(camera.privacy_zones):
@ -171,10 +167,10 @@ def _async_unique_id_to_mac(unique_id: str) -> str:
return unique_id.split("_")[0]
async def set_chime_paired_doorbells(hass: HomeAssistant, call: ServiceCall) -> None:
async def set_chime_paired_doorbells(call: ServiceCall) -> None:
"""Set paired doorbells on chime."""
ref = async_extract_referenced_entity_ids(hass, call)
entity_registry = er.async_get(hass)
ref = async_extract_referenced_entity_ids(call.hass, call)
entity_registry = er.async_get(call.hass)
entity_id = ref.indirectly_referenced.pop()
chime_button = entity_registry.async_get(entity_id)
@ -182,13 +178,13 @@ async def set_chime_paired_doorbells(hass: HomeAssistant, call: ServiceCall) ->
assert chime_button.device_id is not None
chime_mac = _async_unique_id_to_mac(chime_button.unique_id)
instance = _async_get_ufp_instance(hass, chime_button.device_id)
instance = _async_get_ufp_instance(call.hass, chime_button.device_id)
chime = instance.bootstrap.get_device_from_mac(chime_mac)
chime = cast(Chime, chime)
assert chime is not None
call.data = ReadOnlyDict(call.data.get("doorbells") or {})
doorbell_refs = async_extract_referenced_entity_ids(hass, call)
doorbell_refs = async_extract_referenced_entity_ids(call.hass, call)
doorbell_ids: set[str] = set()
for camera_id in doorbell_refs.referenced | doorbell_refs.indirectly_referenced:
doorbell_sensor = entity_registry.async_get(camera_id)
@ -209,31 +205,32 @@ async def set_chime_paired_doorbells(hass: HomeAssistant, call: ServiceCall) ->
await chime.save_device(data_before_changed)
SERVICES = [
(
SERVICE_ADD_DOORBELL_TEXT,
add_doorbell_text,
DOORBELL_TEXT_SCHEMA,
),
(
SERVICE_REMOVE_DOORBELL_TEXT,
remove_doorbell_text,
DOORBELL_TEXT_SCHEMA,
),
(
SERVICE_SET_CHIME_PAIRED,
set_chime_paired_doorbells,
CHIME_PAIRED_SCHEMA,
),
(
SERVICE_REMOVE_PRIVACY_ZONE,
remove_privacy_zone,
REMOVE_PRIVACY_ZONE_SCHEMA,
),
]
def async_setup_services(hass: HomeAssistant) -> None:
"""Set up the global UniFi Protect services."""
services = [
(
SERVICE_ADD_DOORBELL_TEXT,
functools.partial(add_doorbell_text, hass),
DOORBELL_TEXT_SCHEMA,
),
(
SERVICE_REMOVE_DOORBELL_TEXT,
functools.partial(remove_doorbell_text, hass),
DOORBELL_TEXT_SCHEMA,
),
(
SERVICE_SET_CHIME_PAIRED,
functools.partial(set_chime_paired_doorbells, hass),
CHIME_PAIRED_SCHEMA,
),
(
SERVICE_REMOVE_PRIVACY_ZONE,
functools.partial(remove_privacy_zone, hass),
REMOVE_PRIVACY_ZONE_SCHEMA,
),
]
for name, method, schema in services:
if hass.services.has_service(DOMAIN, name):
continue
for name, method, schema in SERVICES:
hass.services.async_register(DOMAIN, name, method, schema=schema)