Add type annotations to Ambient PWS (#52596)

homeassistant/components/ambient_station/__init__.py

@@ -1,4 +1,5 @@
 """Support for Ambient Weather Station Service."""
+from __future__ import annotations
 
 from aioambient import Client
 from aioambient.errors import WebsocketError
@@ -8,6 +9,7 @@ from homeassistant.components.binary_sensor import (
     DOMAIN as BINARY_SENSOR,
 )
 from homeassistant.components.sensor import DOMAIN as SENSOR
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_LOCATION,
     ATTR_NAME,
@@ -30,7 +32,7 @@ from homeassistant.const import (
     SPEED_MILES_PER_HOUR,
     TEMP_FAHRENHEIT,
 )
-from homeassistant.core import callback
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import aiohttp_client, config_validation as cv
 from homeassistant.helpers.dispatcher import (
@@ -292,7 +294,7 @@ SENSOR_TYPES = {
 CONFIG_SCHEMA = cv.deprecated(DOMAIN)
 
 
-async def async_setup_entry(hass, config_entry):
+async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Set up the Ambient PWS as config entry."""
     hass.data.setdefault(DOMAIN, {DATA_CLIENT: {}})
 
@@ -330,7 +332,7 @@ async def async_setup_entry(hass, config_entry):
     return True
 
 
-async def async_unload_entry(hass, config_entry):
+async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Unload an Ambient PWS config entry."""
     ambient = hass.data[DOMAIN][DATA_CLIENT].pop(config_entry.entry_id)
     hass.async_create_task(ambient.ws_disconnect())
@@ -338,7 +340,7 @@ async def async_unload_entry(hass, config_entry):
     return await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS)
 
 
-async def async_migrate_entry(hass, config_entry):
+async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Migrate old entry."""
     version = config_entry.version
 
@@ -362,19 +364,21 @@ async def async_migrate_entry(hass, config_entry):
 class AmbientStation:
     """Define a class to handle the Ambient websocket."""
 
-    def __init__(self, hass, config_entry, client):
+    def __init__(
+        self, hass: HomeAssistant, config_entry: ConfigEntry, client: Client
+    ) -> None:
         """Initialize."""
         self._config_entry = config_entry
         self._entry_setup_complete = False
         self._hass = hass
         self._ws_reconnect_delay = DEFAULT_SOCKET_MIN_RETRY
         self.client = client
-        self.stations = {}
+        self.stations: dict[str, dict] = {}
 
-    async def _attempt_connect(self):
+    async def _attempt_connect(self) -> None:
         """Attempt to connect to the socket (retrying later on fail)."""
 
-        async def connect(timestamp=None):
+        async def connect(timestamp: int | None = None):
             """Connect."""
             await self.client.websocket.connect()
 
@@ -385,14 +389,14 @@ class AmbientStation:
             self._ws_reconnect_delay = min(2 * self._ws_reconnect_delay, 480)
             async_call_later(self._hass, self._ws_reconnect_delay, connect)
 
-    async def ws_connect(self):
+    async def ws_connect(self) -> None:
         """Register handlers and connect to the websocket."""
 
-        def on_connect():
+        def on_connect() -> None:
             """Define a handler to fire when the websocket is connected."""
             LOGGER.info("Connected to websocket")
 
-        def on_data(data):
+        def on_data(data: dict) -> None:
             """Define a handler to fire when the data is received."""
             mac_address = data["macAddress"]
             if data != self.stations[mac_address][ATTR_LAST_DATA]:
@@ -402,11 +406,11 @@ class AmbientStation:
                     self._hass, f"ambient_station_data_update_{mac_address}"
                 )
 
-        def on_disconnect():
+        def on_disconnect() -> None:
             """Define a handler to fire when the websocket is disconnected."""
             LOGGER.info("Disconnected from websocket")
 
-        def on_subscribed(data):
+        def on_subscribed(data: dict) -> None:
             """Define a handler to fire when the subscription is set."""
             for station in data["devices"]:
                 if station["macAddress"] in self.stations:
@@ -447,7 +451,7 @@ class AmbientStation:
 
         await self._attempt_connect()
 
-    async def ws_disconnect(self):
+    async def ws_disconnect(self) -> None:
         """Disconnect from the websocket."""
         await self.client.websocket.disconnect()
 
@@ -456,8 +460,14 @@ class AmbientWeatherEntity(Entity):
     """Define a base Ambient PWS entity."""
 
     def __init__(
-        self, ambient, mac_address, station_name, sensor_type, sensor_name, device_class
-    ):
+        self,
+        ambient: AmbientStation,
+        mac_address: str,
+        station_name: str,
+        sensor_type: str,
+        sensor_name: str,
+        device_class: str | None,
+    ) -> None:
         """Initialize the sensor."""
         self._ambient = ambient
         self._attr_device_class = device_class
@@ -472,11 +482,11 @@ class AmbientWeatherEntity(Entity):
         self._mac_address = mac_address
         self._sensor_type = sensor_type
 
-    async def async_added_to_hass(self):
+    async def async_added_to_hass(self) -> None:
         """Register callbacks."""
 
         @callback
-        def update():
+        def update() -> None:
             """Update the state."""
             if self._sensor_type == TYPE_SOLARRADIATION_LX:
                 self._attr_available = (
@@ -505,6 +515,6 @@ class AmbientWeatherEntity(Entity):
         self.update_from_latest_data()
 
     @callback
-    def update_from_latest_data(self):
+    def update_from_latest_data(self) -> None:
         """Update the entity from the latest data."""
         raise NotImplementedError

homeassistant/components/ambient_station/binary_sensor.py

@@ -1,10 +1,14 @@
 """Support for Ambient Weather Station binary sensors."""
+from __future__ import annotations
+
 from homeassistant.components.binary_sensor import (
     DOMAIN as BINARY_SENSOR,
     BinarySensorEntity,
 )
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_NAME
-from homeassistant.core import callback
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import (
     SENSOR_TYPES,
@@ -27,7 +31,9 @@ from . import (
 from .const import ATTR_LAST_DATA, ATTR_MONITORED_CONDITIONS, DATA_CLIENT, DOMAIN
 
 
-async def async_setup_entry(hass, entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+) -> None:
     """Set up Ambient PWS binary sensors based on a config entry."""
     ambient = hass.data[DOMAIN][DATA_CLIENT][entry.entry_id]
 
@@ -54,7 +60,7 @@ class AmbientWeatherBinarySensor(AmbientWeatherEntity, BinarySensorEntity):
     """Define an Ambient binary sensor."""
 
     @callback
-    def update_from_latest_data(self):
+    def update_from_latest_data(self) -> None:
         """Fetch new state data for the entity."""
         state = self._ambient.stations[self._mac_address][ATTR_LAST_DATA].get(
             self._sensor_type

homeassistant/components/ambient_station/config_flow.py

@@ -1,10 +1,13 @@
 """Config flow to configure the Ambient PWS component."""
+from __future__ import annotations
+
 from aioambient import Client
 from aioambient.errors import AmbientError
 import voluptuous as vol
 
 from homeassistant import config_entries
 from homeassistant.const import CONF_API_KEY
+from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers import aiohttp_client
 
 from .const import CONF_APP_KEY, DOMAIN
@@ -15,13 +18,13 @@ class AmbientStationFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
 
     VERSION = 2
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize the config flow."""
         self.data_schema = vol.Schema(
             {vol.Required(CONF_API_KEY): str, vol.Required(CONF_APP_KEY): str}
         )
 
-    async def _show_form(self, errors=None):
+    async def _show_form(self, errors: dict | None = None) -> FlowResult:
         """Show the form to the user."""
         return self.async_show_form(
             step_id="user",
@@ -29,7 +32,7 @@ class AmbientStationFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
             errors=errors if errors else {},
         )
 
-    async def async_step_user(self, user_input=None):
+    async def async_step_user(self, user_input: dict | None = None) -> FlowResult:
         """Handle the start of the config flow."""
         if not user_input:
             return await self._show_form()

homeassistant/components/ambient_station/sensor.py

@@ -1,18 +1,25 @@
 """Support for Ambient Weather Station sensors."""
+from __future__ import annotations
+
 from homeassistant.components.sensor import DOMAIN as SENSOR, SensorEntity
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_NAME
-from homeassistant.core import callback
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import (
     SENSOR_TYPES,
     TYPE_SOLARRADIATION,
     TYPE_SOLARRADIATION_LX,
+    AmbientStation,
     AmbientWeatherEntity,
 )
 from .const import ATTR_LAST_DATA, ATTR_MONITORED_CONDITIONS, DATA_CLIENT, DOMAIN
 
 
-async def async_setup_entry(hass, entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+) -> None:
     """Set up Ambient PWS sensors based on a config entry."""
     ambient = hass.data[DOMAIN][DATA_CLIENT][entry.entry_id]
 
@@ -41,14 +48,14 @@ class AmbientWeatherSensor(AmbientWeatherEntity, SensorEntity):
 
     def __init__(
         self,
-        ambient,
-        mac_address,
-        station_name,
-        sensor_type,
-        sensor_name,
-        device_class,
-        unit,
-    ):
+        ambient: AmbientStation,
+        mac_address: str,
+        station_name: str,
+        sensor_type: str,
+        sensor_name: str,
+        device_class: str | None,
+        unit: str | None,
+    ) -> None:
         """Initialize the sensor."""
         super().__init__(
             ambient, mac_address, station_name, sensor_type, sensor_name, device_class
@@ -57,7 +64,7 @@ class AmbientWeatherSensor(AmbientWeatherEntity, SensorEntity):
         self._attr_unit_of_measurement = unit
 
     @callback
-    def update_from_latest_data(self):
+    def update_from_latest_data(self) -> None:
         """Fetch new state data for the sensor."""
         if self._sensor_type == TYPE_SOLARRADIATION_LX:
             # If the user requests the solarradiation_lx sensor, use the