Small cleanups to the websocket command phase (#129712)

* Small cleanups to the websocket command phase

- Remove unused argument
- Avoid multiple NamedTuple property lookups

* Update homeassistant/components/websocket_api/http.py

Co-authored-by: epenet <6771947+epenet@users.noreply.github.com>

* Apply suggestions from code review

* touch ups

---------

Co-authored-by: epenet <6771947+epenet@users.noreply.github.com>
This commit is contained in:
J. Nick Koston 2024-11-04 08:33:15 -06:00 committed by GitHub
parent d0c45b1857
commit 7691991a93
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -36,6 +36,8 @@ from .error import Disconnect
from .messages import message_to_json_bytes from .messages import message_to_json_bytes
from .util import describe_request from .util import describe_request
CLOSE_MSG_TYPES = {WSMsgType.CLOSE, WSMsgType.CLOSED, WSMsgType.CLOSING}
if TYPE_CHECKING: if TYPE_CHECKING:
from .connection import ActiveConnection from .connection import ActiveConnection
@ -344,7 +346,7 @@ class WebSocketHandler:
try: try:
connection = await self._async_handle_auth_phase(auth, send_bytes_text) connection = await self._async_handle_auth_phase(auth, send_bytes_text)
self._async_increase_writer_limit(writer) self._async_increase_writer_limit(writer)
await self._async_websocket_command_phase(connection, send_bytes_text) await self._async_websocket_command_phase(connection)
except asyncio.CancelledError: except asyncio.CancelledError:
logger.debug("%s: Connection cancelled", self.description) logger.debug("%s: Connection cancelled", self.description)
raise raise
@ -454,9 +456,7 @@ class WebSocketHandler:
writer._limit = 2**20 # noqa: SLF001 writer._limit = 2**20 # noqa: SLF001
async def _async_websocket_command_phase( async def _async_websocket_command_phase(
self, self, connection: ActiveConnection
connection: ActiveConnection,
send_bytes_text: Callable[[bytes], Coroutine[Any, Any, None]],
) -> None: ) -> None:
"""Handle the command phase of the websocket connection.""" """Handle the command phase of the websocket connection."""
wsock = self._wsock wsock = self._wsock
@ -467,24 +467,26 @@ class WebSocketHandler:
# Command phase # Command phase
while not wsock.closed: while not wsock.closed:
msg = await wsock.receive() msg = await wsock.receive()
msg_type = msg.type
msg_data = msg.data
if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSED, WSMsgType.CLOSING): if msg_type in CLOSE_MSG_TYPES:
break break
if msg.type is WSMsgType.BINARY: if msg_type is WSMsgType.BINARY:
if len(msg.data) < 1: if len(msg_data) < 1:
raise Disconnect("Received invalid binary message.") raise Disconnect("Received invalid binary message.")
handler = msg.data[0] handler = msg_data[0]
payload = msg.data[1:] payload = msg_data[1:]
async_handle_binary(handler, payload) async_handle_binary(handler, payload)
continue continue
if msg.type is not WSMsgType.TEXT: if msg_type is not WSMsgType.TEXT:
raise Disconnect("Received non-Text message.") raise Disconnect("Received non-Text message.")
try: try:
command_msg_data = json_loads(msg.data) command_msg_data = json_loads(msg_data)
except ValueError as ex: except ValueError as ex:
raise Disconnect("Received invalid JSON.") from ex raise Disconnect("Received invalid JSON.") from ex