Tweak type annotations of energy websocket handlers (#119957)

This commit is contained in:
Erik Montnemery 2024-06-19 20:01:02 +02:00 committed by GitHub
parent f0dc39a903
commit 5ee418724b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -4,7 +4,7 @@ from __future__ import annotations
import asyncio
from collections import defaultdict
from collections.abc import Awaitable, Callable
from collections.abc import Callable, Coroutine
from datetime import timedelta
import functools
from itertools import chain
@ -39,7 +39,7 @@ type EnergyWebSocketCommandHandler = Callable[
]
type AsyncEnergyWebSocketCommandHandler = Callable[
[HomeAssistant, websocket_api.ActiveConnection, dict[str, Any], EnergyManager],
Awaitable[None],
Coroutine[Any, Any, None],
]
@ -81,11 +81,10 @@ async def async_get_energy_platforms(
def _ws_with_manager(
func: Any,
) -> websocket_api.WebSocketCommandHandler:
func: AsyncEnergyWebSocketCommandHandler | EnergyWebSocketCommandHandler,
) -> websocket_api.AsyncWebSocketCommandHandler:
"""Decorate a function to pass in a manager."""
@websocket_api.async_response
@functools.wraps(func)
async def with_manager(
hass: HomeAssistant,
@ -107,12 +106,13 @@ def _ws_with_manager(
vol.Required("type"): "energy/get_prefs",
}
)
@websocket_api.async_response
@_ws_with_manager
@callback
def ws_get_prefs(
hass: HomeAssistant,
connection: websocket_api.ActiveConnection,
msg: dict,
msg: dict[str, Any],
manager: EnergyManager,
) -> None:
"""Handle get prefs command."""
@ -131,11 +131,12 @@ def ws_get_prefs(
vol.Optional("device_consumption"): [DEVICE_CONSUMPTION_SCHEMA],
}
)
@websocket_api.async_response
@_ws_with_manager
async def ws_save_prefs(
hass: HomeAssistant,
connection: websocket_api.ActiveConnection,
msg: dict,
msg: dict[str, Any],
manager: EnergyManager,
) -> None:
"""Handle get prefs command."""
@ -187,6 +188,7 @@ async def ws_validate(
vol.Required("type"): "energy/solar_forecast",
}
)
@websocket_api.async_response
@_ws_with_manager
async def ws_solar_forecast(
hass: HomeAssistant,