From 07104a971b4a1ebf8e6d6b828b865732f874d5e5 Mon Sep 17 00:00:00 2001
From: Marc Mueller <30130371+cdce8p@users.noreply.github.com>
Date: Tue, 25 Jun 2024 13:44:13 +0200
Subject: [PATCH] Improve websocket_api schema typing

---
 homeassistant/components/websocket_api/__init__.py   | 6 ++----
 homeassistant/components/websocket_api/decorators.py | 3 ++-
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/homeassistant/components/websocket_api/__init__.py b/homeassistant/components/websocket_api/__init__.py
index d8427bff10e35c..f9bc4396e01b0d 100644
--- a/homeassistant/components/websocket_api/__init__.py
+++ b/homeassistant/components/websocket_api/__init__.py
@@ -4,11 +4,9 @@
 
 from typing import Final, cast
 
-import voluptuous as vol
-
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import config_validation as cv
-from homeassistant.helpers.typing import ConfigType
+from homeassistant.helpers.typing import ConfigType, VolSchemaType
 from homeassistant.loader import bind_hass
 
 from . import commands, connection, const, decorators, http, messages  # noqa: F401
@@ -55,7 +53,7 @@ def async_register_command(
     hass: HomeAssistant,
     command_or_handler: str | const.WebSocketCommandHandler,
     handler: const.WebSocketCommandHandler | None = None,
-    schema: vol.Schema | None = None,
+    schema: VolSchemaType | None = None,
 ) -> None:
     """Register a websocket command."""
     if handler is None:
diff --git a/homeassistant/components/websocket_api/decorators.py b/homeassistant/components/websocket_api/decorators.py
index 5131d02b4d3c8d..b9924bc91d1fe7 100644
--- a/homeassistant/components/websocket_api/decorators.py
+++ b/homeassistant/components/websocket_api/decorators.py
@@ -11,6 +11,7 @@
 from homeassistant.const import HASSIO_USER_NAME
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import Unauthorized
+from homeassistant.helpers.typing import VolDictType
 
 from . import const, messages
 from .connection import ActiveConnection
@@ -130,7 +131,7 @@ def output_error(message_id: str, message: str) -> None:
 
 
 def websocket_command(
-    schema: dict[vol.Marker, Any] | vol.All,
+    schema: VolDictType | vol.All,
 ) -> Callable[[const.WebSocketCommandHandler], const.WebSocketCommandHandler]:
     """Tag a function as a websocket command.