From 4b2e4c8276e7dea480ea32dfee87a629f65a84ec Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Tue, 30 Aug 2022 21:51:21 +0200
Subject: [PATCH] Improve type hints in demo [3/3] (#77186)

---
 homeassistant/components/demo/__init__.py | 89 +++++++++++++----------
 1 file changed, 51 insertions(+), 38 deletions(-)

diff --git a/homeassistant/components/demo/__init__.py b/homeassistant/components/demo/__init__.py
index 8572ddfbbe2..cc80edbd484 100644
--- a/homeassistant/components/demo/__init__.py
+++ b/homeassistant/components/demo/__init__.py
@@ -1,4 +1,6 @@
 """Set up the demo environment that mimics interaction with devices."""
+from __future__ import annotations
+
 import asyncio
 import datetime
 from random import random
@@ -6,7 +8,7 @@ from random import random
 from homeassistant import config_entries, setup
 from homeassistant.components import persistent_notification
 from homeassistant.components.recorder import get_instance
-from homeassistant.components.recorder.models import StatisticMetaData
+from homeassistant.components.recorder.models import StatisticData, StatisticMetaData
 from homeassistant.components.recorder.statistics import (
     async_add_external_statistics,
     get_last_statistics,
@@ -16,9 +18,10 @@ from homeassistant.const import (
     ATTR_ENTITY_ID,
     EVENT_HOMEASSISTANT_START,
     SOUND_PRESSURE_DB,
+    Platform,
 )
 import homeassistant.core as ha
-from homeassistant.core import HomeAssistant
+from homeassistant.core import Event, HomeAssistant
 from homeassistant.helpers.discovery import async_load_platform
 from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue
 from homeassistant.helpers.typing import ConfigType
@@ -27,36 +30,36 @@ import homeassistant.util.dt as dt_util
 DOMAIN = "demo"
 
 COMPONENTS_WITH_CONFIG_ENTRY_DEMO_PLATFORM = [
-    "air_quality",
-    "alarm_control_panel",
-    "binary_sensor",
-    "button",
-    "camera",
-    "climate",
-    "cover",
-    "fan",
-    "humidifier",
-    "light",
-    "lock",
-    "media_player",
-    "number",
-    "select",
-    "sensor",
-    "siren",
-    "switch",
-    "update",
-    "vacuum",
-    "water_heater",
+    Platform.AIR_QUALITY,
+    Platform.ALARM_CONTROL_PANEL,
+    Platform.BINARY_SENSOR,
+    Platform.BUTTON,
+    Platform.CAMERA,
+    Platform.CLIMATE,
+    Platform.COVER,
+    Platform.FAN,
+    Platform.HUMIDIFIER,
+    Platform.LIGHT,
+    Platform.LOCK,
+    Platform.MEDIA_PLAYER,
+    Platform.NUMBER,
+    Platform.SELECT,
+    Platform.SENSOR,
+    Platform.SIREN,
+    Platform.SWITCH,
+    Platform.UPDATE,
+    Platform.VACUUM,
+    Platform.WATER_HEATER,
 ]
 
 COMPONENTS_WITH_DEMO_PLATFORM = [
-    "tts",
-    "stt",
-    "mailbox",
-    "notify",
-    "image_processing",
-    "calendar",
-    "device_tracker",
+    Platform.TTS,
+    Platform.STT,
+    Platform.MAILBOX,
+    Platform.NOTIFY,
+    Platform.IMAGE_PROCESSING,
+    Platform.CALENDAR,
+    Platform.DEVICE_TRACKER,
 ]
 
 
@@ -173,7 +176,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
         title="Example Notification",
     )
 
-    async def demo_start_listener(_event):
+    async def demo_start_listener(_event: Event) -> None:
         """Finish set up."""
         await finish_setup(hass, config)
 
@@ -225,8 +228,10 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
-def _generate_mean_statistics(start, end, init_value, max_diff):
-    statistics = []
+def _generate_mean_statistics(
+    start: datetime.datetime, end: datetime.datetime, init_value: float, max_diff: float
+) -> list[StatisticData]:
+    statistics: list[StatisticData] = []
     mean = init_value
     now = start
     while now < end:
@@ -244,10 +249,16 @@ def _generate_mean_statistics(start, end, init_value, max_diff):
     return statistics
 
 
-async def _insert_sum_statistics(hass, metadata, start, end, max_diff):
-    statistics = []
+async def _insert_sum_statistics(
+    hass: HomeAssistant,
+    metadata: StatisticMetaData,
+    start: datetime.datetime,
+    end: datetime.datetime,
+    max_diff: float,
+):
+    statistics: list[StatisticData] = []
     now = start
-    sum_ = 0
+    sum_ = 0.0
     statistic_id = metadata["statistic_id"]
 
     last_stats = await get_instance(hass).async_add_executor_job(
@@ -349,10 +360,10 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
     return True
 
 
-async def finish_setup(hass, config):
+async def finish_setup(hass: HomeAssistant, config: ConfigType) -> None:
     """Finish set up once demo platforms are set up."""
-    switches = None
-    lights = None
+    switches: list[str] | None = None
+    lights: list[str] | None = None
 
     while not switches and not lights:
         # Not all platforms might be loaded.
@@ -361,6 +372,8 @@ async def finish_setup(hass, config):
         switches = sorted(hass.states.async_entity_ids("switch"))
         lights = sorted(hass.states.async_entity_ids("light"))
 
+    assert switches is not None
+    assert lights is not None
     # Set up scripts
     await setup.async_setup_component(
         hass,