Add type hints to plex data (#89221)

* Add type hints to plex data

* Rename method

homeassistant/components/plex/__init__.py

@@ -35,8 +35,6 @@ from .const import (
     CONF_SERVER_IDENTIFIER,
     DISPATCHERS,
     DOMAIN,
-    GDM_DEBOUNCER,
-    GDM_SCANNER,
     PLATFORMS,
     PLATFORMS_COMPLETED,
     PLEX_SERVER_CONFIG,
@@ -47,6 +45,7 @@ from .const import (
     WEBSOCKETS,
 )
 from .errors import ShouldUpdateConfigEntry
+from .helpers import PlexData, get_plex_data
 from .media_browser import browse_media
 from .server import PlexServer
 from .services import async_setup_services
@@ -62,7 +61,7 @@ def is_plex_media_id(media_content_id):
 
 async def async_browse_media(hass, media_content_type, media_content_id, platform=None):
     """Browse Plex media."""
-    plex_server = next(iter(hass.data[DOMAIN][SERVERS].values()), None)
+    plex_server = next(iter(get_plex_data(hass)[SERVERS].values()), None)
     if not plex_server:
         raise BrowseError("No Plex servers available")
     is_internal = is_internal_request(hass)
@@ -80,22 +79,13 @@ async def async_browse_media(hass, media_content_type, media_content_id, platfor
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up the Plex component."""
-    hass.data.setdefault(
+    gdm = GDM()
-        DOMAIN,
-        {SERVERS: {}, DISPATCHERS: {}, WEBSOCKETS: {}, PLATFORMS_COMPLETED: {}},
-    )
-
-    await async_setup_services(hass)
-
-    hass.http.register_view(PlexImageView())
-
-    gdm = hass.data[DOMAIN][GDM_SCANNER] = GDM()
 
     def gdm_scan():
         _LOGGER.debug("Scanning for GDM clients")
         gdm.scan(scan_for_clients=True)
 
-    hass.data[DOMAIN][GDM_DEBOUNCER] = Debouncer[None](
+    debouncer = Debouncer[None](
         hass,
         _LOGGER,
         cooldown=10,
@@ -103,6 +93,20 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
         function=gdm_scan,
     ).async_call
 
+    hass_data = PlexData(
+        servers={},
+        dispatchers={},
+        websockets={},
+        platforms_completed={},
+        gdm_scanner=gdm,
+        gdm_debouncer=debouncer,
+    )
+    hass.data.setdefault(DOMAIN, hass_data)
+
+    await async_setup_services(hass)
+
+    hass.http.register_view(PlexImageView())
+
     return True
 
 
@@ -161,8 +165,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         "Connected to: %s (%s)", plex_server.friendly_name, plex_server.url_in_use
     )
     server_id = plex_server.machine_identifier
-    hass.data[DOMAIN][SERVERS][server_id] = plex_server
+    hass_data = get_plex_data(hass)
-    hass.data[DOMAIN][PLATFORMS_COMPLETED][server_id] = set()
+    hass_data[SERVERS][server_id] = plex_server
+    hass_data[PLATFORMS_COMPLETED][server_id] = set()
 
     entry.add_update_listener(async_options_updated)
 
@@ -171,8 +176,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         PLEX_UPDATE_PLATFORMS_SIGNAL.format(server_id),
         plex_server.async_update_platforms,
     )
-    hass.data[DOMAIN][DISPATCHERS].setdefault(server_id, [])
+    hass_data[DISPATCHERS].setdefault(server_id, [])
-    hass.data[DOMAIN][DISPATCHERS][server_id].append(unsub)
+    hass_data[DISPATCHERS][server_id].append(unsub)
 
     @callback
     def plex_websocket_callback(msgtype, data, error):
@@ -213,11 +218,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         session=session,
         verify_ssl=verify_ssl,
     )
-    hass.data[DOMAIN][WEBSOCKETS][server_id] = websocket
+    hass_data[WEBSOCKETS][server_id] = websocket
 
     def start_websocket_session(platform):
-        hass.data[DOMAIN][PLATFORMS_COMPLETED][server_id].add(platform)
+        hass_data[PLATFORMS_COMPLETED][server_id].add(platform)
-        if hass.data[DOMAIN][PLATFORMS_COMPLETED][server_id] == PLATFORMS:
+        if hass_data[PLATFORMS_COMPLETED][server_id] == PLATFORMS:
             hass.loop.create_task(websocket.listen())
 
     def close_websocket_session(_):
@@ -226,7 +231,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     unsub = hass.bus.async_listen_once(
         EVENT_HOMEASSISTANT_STOP, close_websocket_session
     )
-    hass.data[DOMAIN][DISPATCHERS][server_id].append(unsub)
+    hass_data[DISPATCHERS][server_id].append(unsub)
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -263,16 +268,17 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     server_id = entry.data[CONF_SERVER_IDENTIFIER]
 
-    websocket = hass.data[DOMAIN][WEBSOCKETS].pop(server_id)
+    hass_data = get_plex_data(hass)
+    websocket = hass_data[WEBSOCKETS].pop(server_id)
     websocket.close()
 
-    dispatchers = hass.data[DOMAIN][DISPATCHERS].pop(server_id)
+    dispatchers = hass_data[DISPATCHERS].pop(server_id)
     for unsub in dispatchers:
         unsub()
 
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
-    hass.data[DOMAIN][SERVERS].pop(server_id)
+    hass_data[SERVERS].pop(server_id)
 
     return unload_ok
 
@@ -281,9 +287,10 @@ async def async_options_updated(hass: HomeAssistant, entry: ConfigEntry) -> None
     """Triggered by config entry options updates."""
     server_id = entry.data[CONF_SERVER_IDENTIFIER]
 
+    hass_data = get_plex_data(hass)
     # Guard incomplete setup during reauth flows
-    if server_id in hass.data[DOMAIN][SERVERS]:
+    if server_id in hass_data[SERVERS]:
-        hass.data[DOMAIN][SERVERS][server_id].options = entry.options
+        hass_data[SERVERS][server_id].options = entry.options
 
 
 @callback

homeassistant/components/plex/config_flow.py

@@ -49,13 +49,13 @@ from .const import (
     DOMAIN,
     MANUAL_SETUP_STRING,
     PLEX_SERVER_CONFIG,
-    SERVERS,
     X_PLEX_DEVICE_NAME,
     X_PLEX_PLATFORM,
     X_PLEX_PRODUCT,
     X_PLEX_VERSION,
 )
 from .errors import NoServersFound, ServerNotSpecified
+from .helpers import get_plex_server
 from .server import PlexServer
 
 HEADER_FRONTEND_BASE = "HA-Frontend-Base"
@@ -360,7 +360,7 @@ class PlexOptionsFlowHandler(config_entries.OptionsFlow):
 
     async def async_step_plex_mp_settings(self, user_input=None):
         """Manage the Plex media_player options."""
-        plex_server = self.hass.data[DOMAIN][SERVERS][self.server_id]
+        plex_server = get_plex_server(self.hass, self.server_id)
 
         if user_input is not None:
             self.options[MP_DOMAIN][CONF_USE_EPISODE_ART] = user_input[

homeassistant/components/plex/const.py

@@ -1,5 +1,6 @@
 """Constants for the Plex component."""
 from datetime import timedelta
+from typing import Final
 
 from homeassistant.const import Platform, __version__
 
@@ -16,14 +17,14 @@ PLEXTV_THROTTLE = 60
 
 CLIENT_SCAN_INTERVAL = timedelta(minutes=10)
 DEBOUNCE_TIMEOUT = 1
-DISPATCHERS = "dispatchers"
+DISPATCHERS: Final = "dispatchers"
-GDM_DEBOUNCER = "gdm_debouncer"
+GDM_DEBOUNCER: Final = "gdm_debouncer"
-GDM_SCANNER = "gdm_scanner"
+GDM_SCANNER: Final = "gdm_scanner"
 PLATFORMS = frozenset([Platform.BUTTON, Platform.MEDIA_PLAYER, Platform.SENSOR])
-PLATFORMS_COMPLETED = "platforms_completed"
+PLATFORMS_COMPLETED: Final = "platforms_completed"
 PLAYER_SOURCE = "player_source"
-SERVERS = "servers"
+SERVERS: Final = "servers"
-WEBSOCKETS = "websockets"
+WEBSOCKETS: Final = "websockets"
 
 PLEX_SERVER_CONFIG = "server_config"
 

homeassistant/components/plex/helpers.py

@@ -1,4 +1,40 @@
 """Helper methods for common Plex integration operations."""
+from __future__ import annotations
+
+from collections.abc import Callable, Coroutine
+from typing import TYPE_CHECKING, Any, TypedDict
+
+from plexapi.gdm import GDM
+from plexwebsocket import PlexWebsocket
+
+from homeassistant.const import Platform
+from homeassistant.core import CALLBACK_TYPE, HomeAssistant
+
+from .const import DOMAIN, SERVERS
+
+if TYPE_CHECKING:
+    from . import PlexServer
+
+
+class PlexData(TypedDict):
+    """Typed description of plex data stored in `hass.data`."""
+
+    servers: dict[str, PlexServer]
+    dispatchers: dict[str, list[CALLBACK_TYPE]]
+    websockets: dict[str, PlexWebsocket]
+    platforms_completed: dict[str, set[Platform]]
+    gdm_scanner: GDM
+    gdm_debouncer: Callable[[], Coroutine[Any, Any, None]]
+
+
+def get_plex_data(hass: HomeAssistant) -> PlexData:
+    """Get typed data from hass.data."""
+    return hass.data[DOMAIN]
+
+
+def get_plex_server(hass: HomeAssistant, server_id: str) -> PlexServer:
+    """Get Plex server from hass.data."""
+    return get_plex_data(hass)[SERVERS][server_id]
 
 
 def pretty_title(media, short_name=False):

homeassistant/components/plex/media_browser.py

@@ -7,7 +7,7 @@ from homeassistant.components.media_player import BrowseError, BrowseMedia, Medi
 
 from .const import DOMAIN, SERVERS
 from .errors import MediaNotFound
-from .helpers import pretty_title
+from .helpers import get_plex_data, get_plex_server, pretty_title
 
 
 class UnknownMediaType(BrowseError):
@@ -42,7 +42,7 @@ def browse_media(  # noqa: C901
     if media_content_id:
         url = URL(media_content_id)
         server_id = url.host
-        plex_server = hass.data[DOMAIN][SERVERS][server_id]
+        plex_server = get_plex_server(hass, server_id)
         if media_content_type == "hub":
             _, hub_location, hub_identifier = url.parts
         elif media_content_type in ["library", "server"] and len(url.parts) > 2:
@@ -294,7 +294,7 @@ def root_payload(hass, is_internal, platform=None):
     """Return root payload for Plex."""
     children = []
 
-    for server_id in hass.data[DOMAIN][SERVERS]:
+    for server_id in get_plex_data(hass)[SERVERS]:
         children.append(
             browse_media(
                 hass,

homeassistant/components/plex/media_player.py

@@ -40,9 +40,9 @@ from .const import (
     PLEX_UPDATE_MEDIA_PLAYER_SESSION_SIGNAL,
     PLEX_UPDATE_MEDIA_PLAYER_SIGNAL,
     PLEX_UPDATE_SENSOR_SIGNAL,
-    SERVERS,
     TRANSIENT_DEVICE_MODELS,
 )
+from .helpers import get_plex_data, get_plex_server
 from .media_browser import browse_media
 from .services import process_plex_payload
 
@@ -85,7 +85,7 @@ async def async_setup_entry(
     unsub = async_dispatcher_connect(
         hass, PLEX_NEW_MP_SIGNAL.format(server_id), async_new_media_players
     )
-    hass.data[DOMAIN][DISPATCHERS][server_id].append(unsub)
+    get_plex_data(hass)[DISPATCHERS][server_id].append(unsub)
     _LOGGER.debug("New entity listener created")
 
 
@@ -94,7 +94,7 @@ def _async_add_entities(hass, registry, async_add_entities, server_id, new_entit
     """Set up Plex media_player entities."""
     _LOGGER.debug("New entities: %s", new_entities)
     entities = []
-    plexserver = hass.data[DOMAIN][SERVERS][server_id]
+    plexserver = get_plex_server(hass, server_id)
     for entity_params in new_entities:
         plex_mp = PlexMediaPlayer(plexserver, **entity_params)
         entities.append(plex_mp)

homeassistant/components/plex/sensor.py

@@ -20,9 +20,8 @@ from .const import (
     NAME_FORMAT,
     PLEX_UPDATE_LIBRARY_SIGNAL,
     PLEX_UPDATE_SENSOR_SIGNAL,
-    SERVERS,
 )
-from .helpers import pretty_title
+from .helpers import get_plex_server, pretty_title
 
 LIBRARY_ATTRIBUTE_TYPES = {
     "artist": ["artist", "album"],
@@ -57,7 +56,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up Plex sensor from a config entry."""
     server_id = config_entry.data[CONF_SERVER_IDENTIFIER]
-    plexserver = hass.data[DOMAIN][SERVERS][server_id]
+    plexserver = get_plex_server(hass, server_id)
     sensors = [PlexSensor(hass, plexserver)]
 
     def create_library_sensors():

homeassistant/components/plex/server.py

@@ -1,4 +1,6 @@
 """Shared class to maintain Plex server instances."""
+from __future__ import annotations
+
 import logging
 import ssl
 import time
@@ -27,7 +29,6 @@ from .const import (
     CONF_USE_EPISODE_ART,
     DEBOUNCE_TIMEOUT,
     DEFAULT_VERIFY_SSL,
-    DOMAIN,
     GDM_DEBOUNCER,
     GDM_SCANNER,
     PLAYER_SOURCE,
@@ -47,6 +48,7 @@ from .errors import (
     ServerNotSpecified,
     ShouldUpdateConfigEntry,
 )
+from .helpers import get_plex_data
 from .media_search import search_media
 from .models import PlexSession
 
@@ -316,7 +318,7 @@ class PlexServer:
         """Update the platform entities."""
         _LOGGER.debug("Updating devices")
 
-        await self.hass.data[DOMAIN][GDM_DEBOUNCER]()
+        await get_plex_data(self.hass)[GDM_DEBOUNCER]()
 
         available_clients = {}
         ignored_clients = set()
@@ -429,7 +431,7 @@ class PlexServer:
 
         def connect_new_clients():
             """Create connections to newly discovered clients."""
-            for gdm_entry in self.hass.data[DOMAIN][GDM_SCANNER].entries:
+            for gdm_entry in get_plex_data(self.hass)[GDM_SCANNER].entries:
                 machine_identifier = gdm_entry["data"]["Resource-Identifier"]
                 if machine_identifier in self._client_device_cache:
                     client = self._client_device_cache[machine_identifier]

homeassistant/components/plex/services.py

@@ -19,6 +19,7 @@ from .const import (
     SERVICE_SCAN_CLIENTS,
 )
 from .errors import MediaNotFound
+from .helpers import get_plex_data
 from .models import PlexMediaSearchResult
 from .server import PlexServer
 
@@ -41,7 +42,7 @@ async def async_setup_services(hass: HomeAssistant) -> None:
             " Service calls will still work for now but the service will be removed in"
             " a future release"
         )
-        for server_id in hass.data[DOMAIN][SERVERS]:
+        for server_id in get_plex_data(hass)[SERVERS]:
             async_dispatcher_send(hass, PLEX_UPDATE_PLATFORMS_SIGNAL.format(server_id))
 
     hass.services.async_register(
@@ -84,7 +85,7 @@ def get_plex_server(
     """Retrieve a configured Plex server by name."""
     if DOMAIN not in hass.data:
         raise HomeAssistantError("Plex integration not configured")
-    servers: dict[str, PlexServer] = hass.data[DOMAIN][SERVERS]
+    servers: dict[str, PlexServer] = get_plex_data(hass)[SERVERS]
     if not servers:
         raise HomeAssistantError("No Plex servers available")
 

homeassistant/components/plex/view.py

@@ -11,7 +11,8 @@ from aiohttp.typedefs import LooseHeaders
 from homeassistant.components.http import KEY_AUTHENTICATED, HomeAssistantView
 from homeassistant.components.media_player import async_fetch_image
 
-from .const import DOMAIN, SERVERS
+from .const import SERVERS
+from .helpers import get_plex_data
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -33,7 +34,7 @@ class PlexImageView(HomeAssistantView):
             return web.Response(status=HTTPStatus.UNAUTHORIZED)
 
         hass = request.app["hass"]
-        if (server := hass.data[DOMAIN][SERVERS].get(server_id)) is None:
+        if (server := get_plex_data(hass)[SERVERS].get(server_id)) is None:
             return web.Response(status=HTTPStatus.NOT_FOUND)
 
         if (image_url := server.thumbnail_cache.get(media_content_id)) is None: