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