From d47faf6a51ac071bf5db23a45bcb31d8362c91d3 Mon Sep 17 00:00:00 2001 From: Bao Trinh Date: Wed, 6 Aug 2025 21:06:34 -0500 Subject: [PATCH 1/2] chore: improve type hints --- src/hatch_rest_api/aws_http.py | 9 +- src/hatch_rest_api/callbacks.py | 9 +- src/hatch_rest_api/const.py | 8 +- src/hatch_rest_api/contentful.py | 7 +- src/hatch_rest_api/hatch.py | 72 +++++++--- src/hatch_rest_api/rest_iot.py | 64 +++++---- src/hatch_rest_api/rest_mini.py | 23 +-- src/hatch_rest_api/rest_plus.py | 39 ++--- src/hatch_rest_api/restore_iot.py | 39 ++--- src/hatch_rest_api/restore_v5.py | 57 ++++---- .../shadow_client_subscriber.py | 21 +-- src/hatch_rest_api/stub.py | 4 +- src/hatch_rest_api/types.py | 134 +++++++++++++++--- src/hatch_rest_api/util.py | 24 ++-- src/hatch_rest_api/util_bootstrap.py | 45 ++++-- src/hatch_rest_api/util_http.py | 11 +- 16 files changed, 383 insertions(+), 183 deletions(-) diff --git a/src/hatch_rest_api/aws_http.py b/src/hatch_rest_api/aws_http.py index 8e34c16..b56d052 100644 --- a/src/hatch_rest_api/aws_http.py +++ b/src/hatch_rest_api/aws_http.py @@ -3,6 +3,7 @@ from aiohttp import ClientSession, ClientResponse import json +from .types import AwsIotCredentialsResponse from .util_http import request_with_logging _LOGGER = logging.getLogger(__name__) @@ -11,22 +12,22 @@ class AwsHttp: - def __init__(self, client_session: ClientSession = None): + def __init__(self, client_session: ClientSession | None = None): if client_session is None: self.api_session = ClientSession(raise_for_status=True) else: self.api_session = client_session - async def cleanup_client_session(self): + async def cleanup_client_session(self) -> None: await self.api_session.close() @request_with_logging async def _post_request_with_logging_and_errors_raised( - self, url: str, json_body: dict, headers: dict = None + self, url: str, json_body: dict, headers: dict | None = None ) -> ClientResponse: return await self.api_session.post(url=url, json=json_body, headers=headers) - async def aws_credentials(self, region: str, identityId: str, aws_token: str): + async def aws_credentials(self, region: str, identityId: str, aws_token: str) -> AwsIotCredentialsResponse: url = f"https://cognito-identity.{region}.amazonaws.com" json_body = { "IdentityId": identityId, diff --git a/src/hatch_rest_api/callbacks.py b/src/hatch_rest_api/callbacks.py index aca7544..7916548 100644 --- a/src/hatch_rest_api/callbacks.py +++ b/src/hatch_rest_api/callbacks.py @@ -1,18 +1,19 @@ +from collections.abc import Callable import logging _LOGGER = logging.getLogger(__name__) class CallbacksMixin: - def _setup_callbacks(self): - self._callbacks = set() + def _setup_callbacks(self) -> None: + self._callbacks: set[Callable[[], None]] = set() - def register_callback(self, callback) -> None: + def register_callback(self, callback: Callable[[], None]) -> None: if not hasattr(self, "_callbacks"): self._setup_callbacks() self._callbacks.add(callback) - def remove_callback(self, callback) -> None: + def remove_callback(self, callback: Callable[[], None]) -> None: if not hasattr(self, "_callbacks"): self._setup_callbacks() self._callbacks.discard(callback) diff --git a/src/hatch_rest_api/const.py b/src/hatch_rest_api/const.py index a986ce8..054cb70 100644 --- a/src/hatch_rest_api/const.py +++ b/src/hatch_rest_api/const.py @@ -12,7 +12,7 @@ NO_SOUND_ID = 19998 -class RestMiniAudioTrack(Enum): +class RestMiniAudioTrack(int, Enum): NONE = 0 Heartbeat = 10124 Water = 10125 @@ -24,7 +24,7 @@ class RestMiniAudioTrack(Enum): Birds = 10131 -class RestPlusAudioTrack(Enum): +class RestPlusAudioTrack(int, Enum): NONE = 0 Stream = 2 PinkNoise = 3 @@ -39,7 +39,7 @@ class RestPlusAudioTrack(Enum): RockABye = 14 -class RIoTAudioTrack(Enum): +class RIoTAudioTrack(int, Enum): NONE = NO_SOUND_ID BrownNoise = 10200 WhiteNoise = 10137 @@ -62,7 +62,7 @@ class RIoTAudioTrack(Enum): RockABye = 10194 @classmethod - def sound_url_map(cls): + def sound_url_map(cls) -> dict[int, str]: """ Hard-coded list, as some of these values are not returned by the 'sounds' API. These were found from manually browsing the app and playing each song, collecting the necessary values (name, id, url) from the Home Assistant debug logs for the `ha_hatch` custom component integration. diff --git a/src/hatch_rest_api/contentful.py b/src/hatch_rest_api/contentful.py index 1d528e8..4604344 100644 --- a/src/hatch_rest_api/contentful.py +++ b/src/hatch_rest_api/contentful.py @@ -4,6 +4,7 @@ from aiohttp import ClientError, ClientResponse, ClientSession from .errors import RateError +from .types import JsonType _LOGGER = logging.getLogger(__name__) @@ -15,13 +16,13 @@ class Contentful: - def __init__(self, client_session: ClientSession = None): + def __init__(self, client_session: ClientSession | None = None): self.api_session = client_session or ClientSession() - async def cleanup_client_session(self): + async def cleanup_client_session(self) -> None: await self.api_session.close() - async def graphql_query(self, query, auth_token=None, max_retries=3, **variables): + async def graphql_query(self, query: str, auth_token: str | None = None, max_retries: int = 3, **variables: JsonType) -> dict[str, JsonType]: retry_count = 0 while True: try: diff --git a/src/hatch_rest_api/hatch.py b/src/hatch_rest_api/hatch.py index 9f24257..c33c040 100644 --- a/src/hatch_rest_api/hatch.py +++ b/src/hatch_rest_api/hatch.py @@ -1,19 +1,39 @@ import asyncio +from collections.abc import Awaitable, Callable, Mapping, Sequence import logging +from typing import Literal, TypedDict, overload from aiohttp import ClientError, ClientResponse, ClientSession, __version__ from aiohttp.hdrs import USER_AGENT from .errors import AuthError, RateError +from .types import ( + IotDeviceInfo, + IotTokenResponse, + JsonType, + LoginResponse, + Product, + RestIotRoutine, + SimpleSoundContent, +) from .util_http import request_with_logging +type ContentType = Literal["sound", "color", "windDown"] + + +class ContentResponse[T: SimpleSoundContent | Mapping[str, JsonType]](TypedDict): + contentItems: list[T] + + _LOGGER = logging.getLogger(__name__) API_URL: str = "https://data.hatchbaby.com/" -def request_with_logging_and_errors(func): - async def request_with_logging_wrapper(*args, **kwargs): +def request_with_logging_and_errors[**P, T: ClientResponse]( + func: Callable[P, Awaitable[T]], +) -> Callable[P, Awaitable[T]]: + async def request_with_logging_wrapper(*args: P.args, **kwargs: P.kwargs) -> T: response = await func(*args, **kwargs) if response.status == 429: @@ -45,20 +65,20 @@ async def request_with_logging_wrapper(*args, **kwargs): class Hatch: - def __init__(self, client_session: ClientSession = None): + def __init__(self, client_session: ClientSession | None = None): if client_session is None: self.api_session = ClientSession(raise_for_status=True) else: self.api_session = client_session _LOGGER.debug(f"api_session_version: {__version__}") - async def cleanup_client_session(self): + async def cleanup_client_session(self) -> None: await self.api_session.close() @request_with_logging_and_errors @request_with_logging async def _post_request_with_logging_and_errors_raised( - self, url: str, json_body: dict, auth_token: str = None + self, url: str, json_body: dict, auth_token: str | None = None ) -> ClientResponse: headers = {USER_AGENT: "hatch_rest_api"} if auth_token is not None: @@ -68,7 +88,7 @@ async def _post_request_with_logging_and_errors_raised( @request_with_logging @request_with_logging_and_errors async def _get_request_with_logging_and_errors_raised( - self, url: str, auth_token: str = None, params: dict = None + self, url: str, auth_token: str | None = None, params: dict | None = None ) -> ClientResponse: headers = {USER_AGENT: "hatch_rest_api"} if auth_token is not None: @@ -86,10 +106,10 @@ async def login(self, email: str, password: str) -> str: url=url, json_body=json_body ) ) - response_json = await response.json() + response_json: LoginResponse = await response.json() return response_json["token"] - async def member(self, auth_token: str): + async def member(self, auth_token: str) -> dict[str, JsonType]: url = API_URL + "service/app/v2/member" response: ClientResponse = ( await self._get_request_with_logging_and_errors_raised( @@ -99,7 +119,7 @@ async def member(self, auth_token: str): response_json = await response.json() return response_json["payload"] - async def iot_devices(self, auth_token: str): + async def iot_devices(self, auth_token: str) -> list[IotDeviceInfo]: url = API_URL + "service/app/iotDevice/v2/fetch" params = { "iotProducts": [ @@ -119,7 +139,7 @@ async def iot_devices(self, auth_token: str): response_json = await response.json() return response_json["payload"] - async def token(self, auth_token: str): + async def token(self, auth_token: str) -> IotTokenResponse: url = API_URL + "service/app/restPlus/token/v1/fetch" response: ClientResponse = ( await self._get_request_with_logging_and_errors_raised( @@ -129,7 +149,7 @@ async def token(self, auth_token: str): response_json = await response.json() return response_json["payload"] - async def favorites(self, auth_token: str, mac: str): + async def favorites(self, auth_token: str, mac: str) -> list[RestIotRoutine]: url = API_URL + "service/app/routine/v2/fetch" params = {"macAddress": mac} response: ClientResponse = ( @@ -138,11 +158,11 @@ async def favorites(self, auth_token: str, mac: str): ) ) response_json = await response.json() - favorites = response_json["payload"] + favorites: list[RestIotRoutine] = response_json["payload"] favorites.sort(key=lambda x: x.get("displayOrder", float("inf"))) return favorites - async def routines(self, auth_token: str, mac: str): + async def routines(self, auth_token: str, mac: str) -> list[RestIotRoutine]: url = API_URL + "service/app/routine/v2/fetch" params = {"macAddress": mac, "types": "routine"} response: ClientResponse = ( @@ -151,13 +171,33 @@ async def routines(self, auth_token: str, mac: str): ) ) response_json = await response.json() - routines = response_json["payload"] + routines: list[RestIotRoutine] = response_json["payload"] routines.sort(key=lambda x: x.get("displayOrder", float("inf"))) return routines + @overload + async def content( + self, + auth_token: str, + product: Product, + content: Sequence[Literal["sound"]], + max_retries: int = 3, + ) -> ContentResponse[SimpleSoundContent]: ... + @overload + async def content( + self, + auth_token: str, + product: Product, + content: Sequence[Literal["color", "windDown"]], + max_retries: int = 3, + ) -> ContentResponse[Mapping[str, JsonType]]: ... async def content( - self, auth_token: str, product: str, content: list, max_retries: int = 3 - ): + self, + auth_token: str, + product: Product, + content: Sequence[ContentType], + max_retries: int = 3, + ) -> ContentResponse[Mapping[str, JsonType]] | ContentResponse[SimpleSoundContent]: # content options are ["sound", "color", "windDown"] url = API_URL + "service/app/content/v1/fetchByProduct" params = {"product": product, "contentTypes": content} diff --git a/src/hatch_rest_api/rest_iot.py b/src/hatch_rest_api/rest_iot.py index 29ac2f6..b45393f 100644 --- a/src/hatch_rest_api/rest_iot.py +++ b/src/hatch_rest_api/rest_iot.py @@ -1,7 +1,6 @@ import contextlib import logging -from .types import SoundContent, SimpleSoundContent from .util import ( convert_to_percentage, safely_get_json_value, @@ -18,20 +17,27 @@ RIoTAudioTrack, ) from .shadow_client_subscriber import ShadowClientSubscriberMixin +from .types import ( + IotSoundUntil, + JsonType, + RestChargingStatus, + SimpleSoundContent, + SoundContent, +) _LOGGER = logging.getLogger(__name__) class RestIot(ShadowClientSubscriberMixin): - audio_track: RIoTAudioTrack = None - firmware_version: str = None + audio_track: RIoTAudioTrack | None = None + firmware_version: str | None = None volume: int = 0 is_online: bool = False - current_playing: str = "none" + current_playing: str | None = "none" current_id: int = 0 current_step: int = 0 - battery_level: int = None + battery_level: int | None = None color_id: int = NO_COLOR_ID sound_id: int = NO_SOUND_ID red: int = 0 @@ -39,13 +45,13 @@ class RestIot(ShadowClientSubscriberMixin): blue: int = 0 white: int = 0 brightness: int = 0 - charging_status: int = None # Expected values: 0= Not Charging, 3= Charging, plugged in, 5= Charging, on base - clock: int = None - flags: int = None + charging_status: RestChargingStatus | None = None # Expected values: 0= Not Charging, 3= Charging, plugged in, 5= Charging, on base + clock: int | None = None + flags: int | None = None toddler_lock: bool = False - toddler_lock_mode: str = None + toddler_lock_mode: str | None = None - def _update_local_state(self, state): + def _update_local_state(self, state: dict[str, JsonType]) -> None: _LOGGER.debug(f"update local state: {self.device_name}, {state}") if safely_get_json_value(state, "deviceInfo.f") is not None: self.firmware_version = safely_get_json_value(state, "deviceInfo.f") @@ -105,7 +111,7 @@ def _update_local_state(self, state): _LOGGER.debug(f"new state:{self}") self.publish_updates() - def __repr__(self): + def __repr__(self) -> dict[str, JsonType]: return { "device_name": self.device_name, "thing_name": self.thing_name, @@ -135,30 +141,30 @@ def __repr__(self): "toddler_lock_mode": self.toddler_lock_mode, } - def __str__(self): + def __str__(self) -> str: return f"{self.__repr__()}" @property - def is_on(self): + def is_on(self) -> bool: return self.is_light_on or self.is_playing @property - def is_light_on(self): + def is_light_on(self) -> bool: return self.color_id != NO_COLOR_ID and self.color_id != 0 @property - def is_playing(self): + def is_playing(self) -> bool: return self.sound_id != NO_SOUND_ID @property - def is_clock_on(self): + def is_clock_on(self) -> bool: return self.flags is not None and self.flags & RIOT_FLAGS_CLOCK_ON @property - def is_clock_24h(self): + def is_clock_24h(self) -> bool: return self.flags is not None and self.flags & RIOT_FLAGS_CLOCK_24_HOUR - def favorite_names(self, active_only: bool = True): + def favorite_names(self, active_only: bool = True) -> list[str]: names = [] for favorite in self.favorites: if active_only and favorite["active"]: @@ -167,17 +173,17 @@ def favorite_names(self, active_only: bool = True): names.append(f"{favorite['name']}-{favorite['id']}") return names - def set_volume(self, percentage: int): + def set_volume(self, percentage: float) -> None: _LOGGER.debug(f"Setting volume: {percentage}") self._update({"current": {"sound": {"v": convert_from_percentage(percentage)}}}) # Expected string value for mode is "never" or "always". The API also supports "custom" for defining a time range - def set_toddler_lock(self, on: bool): + def set_toddler_lock(self, on: bool) -> None: _LOGGER.debug(f"Setting Toddler On Lock: {on}") mode = "always" if on else "never" self._update({"toddlerLock": {"turnOnMode": mode}}) - def set_clock(self, brightness: int = 0): + def set_clock(self, brightness: int = 0) -> None: _LOGGER.debug(f"Setting clock on: {brightness}") self._update( { @@ -188,17 +194,17 @@ def set_clock(self, brightness: int = 0): } ) - def turn_clock_off(self): + def turn_clock_off(self) -> None: _LOGGER.debug("Turn off clock") self._update({"clock": {"flags": self.flags ^ RIOT_FLAGS_CLOCK_ON, "i": 655}}) # favorite_name_id is expected to be a string of name-id since name alone isn't unique - def set_favorite(self, favorite_name_id: str): + def set_favorite(self, favorite_name_id: str) -> None: _LOGGER.debug(f"Setting favorite: {favorite_name_id}") fav_id = int(favorite_name_id.rsplit("-", 1)[1]) self._update({"current": {"srId": fav_id, "step": 1, "playing": "routine"}}) - def set_audio_track(self, audio_track: RIoTAudioTrack): + def set_audio_track(self, audio_track: RIoTAudioTrack) -> None: _LOGGER.debug(f"Setting audio track: {audio_track}") if audio_track == RIoTAudioTrack.NONE: self.turn_off() @@ -217,7 +223,7 @@ def set_audio_track(self, audio_track: RIoTAudioTrack): "until": "indefinite", }}}) - def set_sound(self, sound_or_id_or_title: SoundContent | SimpleSoundContent | str | int | None, duration: int = 0, until="indefinite"): + def set_sound(self, sound_or_id_or_title: SoundContent | SimpleSoundContent | str | int | None, duration: int = 0, until: IotSoundUntil = "indefinite") -> None: """ Set a sound by passing SoundContent item from self.sounds, id, or title. """ @@ -253,7 +259,7 @@ def set_sound(self, sound_or_id_or_title: SoundContent | SimpleSoundContent | st } ) - def set_sound_url(self, sound_url: str = 'http://codeskulptor-demos.commondatastorage.googleapis.com/GalaxyInvaders/theme_01.mp3'): + def set_sound_url(self, sound_url: str = 'http://codeskulptor-demos.commondatastorage.googleapis.com/GalaxyInvaders/theme_01.mp3') -> None: """ appears to work with some but not all public wav and mp3 urls @@ -269,11 +275,11 @@ def set_sound_url(self, sound_url: str = 'http://codeskulptor-demos.commondatast } ) - def turn_off(self): + def turn_off(self) -> None: _LOGGER.debug("Turning off sound") self._update({"current": {"srId": 0, "step": 0, "playing": "none"}}) - def turn_light_off(self): + def turn_light_off(self) -> None: _LOGGER.debug("Turning light off") # if favorite is playing then light can be turned off without turning off sound if self.current_playing == "routine": @@ -308,7 +314,7 @@ def turn_light_off(self): def set_color( self, red: int, green: int, blue: int, white: int = 0, brightness: int = 0 - ): + ) -> None: new_color_id = CUSTOM_COLOR_ID _LOGGER.debug( f"red: {red} green: {green} blue: {blue} brightness: {brightness} white: {white}" diff --git a/src/hatch_rest_api/rest_mini.py b/src/hatch_rest_api/rest_mini.py index e574641..070fa04 100644 --- a/src/hatch_rest_api/rest_mini.py +++ b/src/hatch_rest_api/rest_mini.py @@ -1,22 +1,23 @@ import logging -from .util import convert_to_percentage, safely_get_json_value, convert_from_percentage -from .shadow_client_subscriber import ShadowClientSubscriberMixin from .const import RestMiniAudioTrack +from .shadow_client_subscriber import ShadowClientSubscriberMixin +from .types import JsonType +from .util import convert_from_percentage, convert_to_percentage, safely_get_json_value _LOGGER = logging.getLogger(__name__) class RestMini(ShadowClientSubscriberMixin): - firmware_version: str = None - is_online: bool = None + firmware_version: str | None = None + is_online: bool | None = None - audio_track: RestMiniAudioTrack = None - volume: int = None + audio_track: RestMiniAudioTrack | None = None + volume: int | None = None current_playing: str = "none" - def __repr__(self): + def __repr__(self) -> dict[str, JsonType]: return { "device_name": self.device_name, "thing_name": self.thing_name, @@ -29,10 +30,10 @@ def __repr__(self): "document_version": self.document_version, } - def __str__(self): + def __str__(self) -> str: return f"{self.__repr__()}" - def _update_local_state(self, state): + def _update_local_state(self, state: dict[str, JsonType]) -> None: _LOGGER.debug(f"update local state: {self.device_name}, {state}") if safely_get_json_value(state, "connected") is not None: self.is_online = safely_get_json_value(state, "connected") @@ -54,7 +55,7 @@ def _update_local_state(self, state): def is_playing(self) -> bool: return self.current_playing != "none" - def set_volume(self, percentage: int): + def set_volume(self, percentage: float) -> None: self._update( { "current": { @@ -65,7 +66,7 @@ def set_volume(self, percentage: int): } ) - def set_audio_track(self, audio_track: RestMiniAudioTrack): + def set_audio_track(self, audio_track: RestMiniAudioTrack) -> None: if audio_track == RestMiniAudioTrack.NONE: self._update( { diff --git a/src/hatch_rest_api/rest_plus.py b/src/hatch_rest_api/rest_plus.py index 252b776..d8ed439 100644 --- a/src/hatch_rest_api/rest_plus.py +++ b/src/hatch_rest_api/rest_plus.py @@ -8,29 +8,30 @@ convert_to_hex, ) from .shadow_client_subscriber import ShadowClientSubscriberMixin +from .types import JsonType from .const import RestPlusAudioTrack _LOGGER = logging.getLogger(__name__) class RestPlus(ShadowClientSubscriberMixin): - firmware_version: str = None - audio_track: RestPlusAudioTrack = None + firmware_version: str | None = None + audio_track: RestPlusAudioTrack | None = None volume: int = 0 - is_on: bool = None - battery_level: int = None - is_online: bool = None + is_on: bool | None = None + battery_level: int | None = None + is_online: bool | None = None - red: int = None - green: int = None - blue: int = None - brightness: int = None + red: int | None = None + green: int | None = None + blue: int | None = None + brightness: int | None = None - color_random: bool = None - color_white: bool = None + color_random: bool | None = None + color_white: bool | None = None - def _update_local_state(self, state): + def _update_local_state(self, state: dict[str, JsonType]) -> None: _LOGGER.debug(f"update local state: {self.device_name}, {state}") if safely_get_json_value(state, "deviceInfo.f") is not None: self.firmware_version = safely_get_json_value(state, "deviceInfo.f") @@ -74,10 +75,10 @@ def _update_local_state(self, state): self.publish_updates() @property - def is_playing(self): + def is_playing(self) -> bool: return self.is_on and self.audio_track != RestPlusAudioTrack.NONE - def __repr__(self): + def __repr__(self) -> dict[str, JsonType]: return { "device_name": self.device_name, "thing_name": self.thing_name, @@ -96,10 +97,10 @@ def __repr__(self): "document_version": self.document_version, } - def __str__(self): + def __str__(self) -> str: return f"{self.__repr__()}" - def set_volume(self, percentage: int): + def set_volume(self, percentage: float) -> None: self._update( { "a": { @@ -108,7 +109,7 @@ def set_volume(self, percentage: int): } ) - def set_audio_track(self, audio_track: RestPlusAudioTrack): + def set_audio_track(self, audio_track: RestPlusAudioTrack) -> None: self._update( { "a": { @@ -117,10 +118,10 @@ def set_audio_track(self, audio_track: RestPlusAudioTrack): } ) - def set_on(self, on: bool): + def set_on(self, on: bool) -> None: self._update({"isPowered": on}) - def set_color(self, red: int, green: int, blue: int, brightness: int, random: bool = False): + def set_color(self, red: int, green: int, blue: int, brightness: int, random: bool = False) -> None: self._update( { "c": { diff --git a/src/hatch_rest_api/restore_iot.py b/src/hatch_rest_api/restore_iot.py index 9ebd762..f4feff7 100644 --- a/src/hatch_rest_api/restore_iot.py +++ b/src/hatch_rest_api/restore_iot.py @@ -15,13 +15,14 @@ CUSTOM_COLOR_ID, ) from .shadow_client_subscriber import ShadowClientSubscriberMixin +from .types import JsonType, RestChargingStatus _LOGGER = logging.getLogger(__name__) class RestoreIot(ShadowClientSubscriberMixin): - audio_track: str = None - firmware_version: str = None + audio_track: str | None = None + firmware_version: str | None = None volume: int = 0 is_online: bool = False @@ -34,11 +35,11 @@ class RestoreIot(ShadowClientSubscriberMixin): blue: int = 0 white: int = 0 brightness: int = 0 - charging_status: int = None # Expected values: 0= Not Charging, 3= Charging, plugged in, 5= Charging, on base + charging_status: RestChargingStatus | None = None # Expected values: 0= Not Charging, 3= Charging, plugged in, 5= Charging, on base clock: int = 0 flags: int = 0 - def _update_local_state(self, state): + def _update_local_state(self, state: dict[str, JsonType]) -> None: _LOGGER.debug(f"update local state: {self.device_name}, {state}") if safely_get_json_value(state, "deviceInfo.f") is not None: self.firmware_version = safely_get_json_value(state, "deviceInfo.f") @@ -85,7 +86,7 @@ def _update_local_state(self, state): _LOGGER.debug(f"new state:{self}") self.publish_updates() - def __repr__(self): + def __repr__(self) -> dict[str, JsonType]: return { "device_name": self.device_name, "thing_name": self.thing_name, @@ -108,34 +109,34 @@ def __repr__(self): "is_clock_24h": self.is_clock_24h, } - def __str__(self): + def __str__(self) -> str: return f"{self.__repr__()}" @property - def is_on(self): + def is_on(self) -> bool: return self.is_light_on or self.is_playing @property - def is_light_on(self): + def is_light_on(self) -> bool: return self.color_id != NO_COLOR_ID and self.color_id != 0 @property - def is_playing(self): + def is_playing(self) -> bool: return self.sound_id != NO_SOUND_ID @property - def is_clock_on(self): + def is_clock_on(self) -> bool: return self.flags is not None and self.flags & RIOT_FLAGS_CLOCK_ON @property - def is_clock_24h(self): + def is_clock_24h(self) -> bool: return self.flags is not None and self.flags & RIOT_FLAGS_CLOCK_24_HOUR - def set_volume(self, percentage: int): + def set_volume(self, percentage: float) -> None: _LOGGER.debug(f"Setting volume: {percentage}") self._update({"current": {"sound": {"v": convert_from_percentage(percentage)}}}) - def favorite_names(self, active_only: bool = True): + def favorite_names(self, active_only: bool = True) -> list[str]: names = [] for favorite in self.favorites: if active_only and favorite["active"]: @@ -144,27 +145,27 @@ def favorite_names(self, active_only: bool = True): names.append(f"{favorite['name']}-{favorite['id']}") return names - def set_clock(self, brightness: int = 0): + def set_clock(self, brightness: int = 0) -> None: _LOGGER.debug(f"Setting clock on: {brightness}") self._update( {"clock": {"flags": self.flags | RIOT_FLAGS_CLOCK_ON, "i": convert_from_percentage(brightness)}} ) - def turn_clock_off(self): + def turn_clock_off(self) -> None: _LOGGER.debug("Turn off clock") self._update({"clock": {"flags": self.flags ^ RIOT_FLAGS_CLOCK_ON, "i": 655}}) # favorite_name_id is expected to be a string of name-id since name alone isn't unique - def set_favorite(self, favorite_name_id: str): + def set_favorite(self, favorite_name_id: str) -> None: _LOGGER.debug(f"Setting favorite: {favorite_name_id}") fav_id = int(favorite_name_id.rsplit("-", 1)[1]) self._update({"current": {"srId": fav_id, "step": 1, "playing": "routine"}}) - def turn_off(self): + def turn_off(self) -> None: _LOGGER.debug("Turning off sound") self._update({"current": {"srId": 0, "step": 0, "playing": "none"}}) - def turn_light_off(self): + def turn_light_off(self) -> None: _LOGGER.debug("Turning light off") # if favorite is playing then light can be turned off without turning off sound if self.current_playing == "routine": @@ -199,7 +200,7 @@ def turn_light_off(self): def set_color( self, red: int, green: int, blue: int, white: int = 0, brightness: int = 0 - ): + ) -> None: new_color_id = CUSTOM_COLOR_ID _LOGGER.debug( f"red: {red} green: {green} blue: {blue} brightness: {brightness} white: {white}" diff --git a/src/hatch_rest_api/restore_v5.py b/src/hatch_rest_api/restore_v5.py index da68d73..4faac10 100644 --- a/src/hatch_rest_api/restore_v5.py +++ b/src/hatch_rest_api/restore_v5.py @@ -1,6 +1,6 @@ import logging +from typing import cast -from .types import SoundContent, SimpleSoundContent from .util import ( convert_to_percentage, safely_get_json_value, @@ -16,12 +16,13 @@ CUSTOM_COLOR_ID, ) from .shadow_client_subscriber import ShadowClientSubscriberMixin +from .types import IotSoundUntil, JsonType, SimpleSoundContent, SoundContent _LOGGER = logging.getLogger(__name__) class RestoreV5(ShadowClientSubscriberMixin): - firmware_version: str = None + firmware_version: str | None = None volume: int = 0 is_online: bool = False @@ -39,7 +40,7 @@ class RestoreV5(ShadowClientSubscriberMixin): clock_daytime: int = 0 flags: int = 0 - def _update_local_state(self, state): + def _update_local_state(self, state: dict[str, JsonType]) -> None: _LOGGER.debug(f"update local state: {self.device_name}, {state}") if safely_get_json_value(state, "deviceInfo.f") is not None: self.firmware_version = safely_get_json_value(state, "deviceInfo.f") @@ -87,7 +88,7 @@ def _update_local_state(self, state): _LOGGER.debug(f"new state:{self}") self.publish_updates() - def __repr__(self): + def __repr__(self) -> dict[str, JsonType]: return { "device_name": self.device_name, "thing_name": self.thing_name, @@ -113,38 +114,38 @@ def __repr__(self): "is_clock_24h": self.is_clock_24h, } - def __str__(self): + def __str__(self) -> str: return f"{self.__repr__()}" @property - def is_on(self): + def is_on(self) -> bool: return self.is_light_on or self.is_playing @property - def is_light_on(self): + def is_light_on(self) -> bool: return self.color_id != NO_COLOR_ID and self.color_id != 0 @property - def is_playing(self): + def is_playing(self) -> bool: return self.sound_id != NO_SOUND_ID @property - def is_clock_on(self): + def is_clock_on(self) -> bool: return self.flags is not None and self.flags & RIOT_FLAGS_CLOCK_ON @property - def is_clock_24h(self): + def is_clock_24h(self) -> bool: return self.flags is not None and self.flags & RIOT_FLAGS_CLOCK_24_HOUR @property def clock(self) -> int: return self.clock_daytime - def set_volume(self, percentage: int): + def set_volume(self, percentage: float) -> None: _LOGGER.debug(f"Setting volume: {percentage}") self._update({"current": {"sound": {"v": convert_from_percentage(percentage)}}}) - def favorite_names(self, active_only: bool = True): + def favorite_names(self, active_only: bool = True) -> list[str]: names = [] for favorite in self.favorites: if active_only and favorite["active"]: @@ -153,7 +154,7 @@ def favorite_names(self, active_only: bool = True): names.append(f"{favorite['name']}-{favorite['id']}") return names - def set_clock(self, daytime_brightness: int|None = None, nighttime_brightness: int|None = None): + def set_clock(self, daytime_brightness: int | None = None, nighttime_brightness: int | None = None) -> None: if daytime_brightness is None: daytime_brightness = self.clock_daytime if nighttime_brightness is None: @@ -163,20 +164,23 @@ def set_clock(self, daytime_brightness: int|None = None, nighttime_brightness: i {"clock": {"flags": self.flags | RIOT_FLAGS_CLOCK_ON, "i": pack_dual_percentages(nighttime_brightness, daytime_brightness)}} ) - def turn_clock_off(self): + def turn_clock_off(self) -> None: _LOGGER.debug("Turn off clock") self._update({"clock": {"flags": self.flags ^ RIOT_FLAGS_CLOCK_ON, "i": 655}}) # favorite_name_id is expected to be a string of name-id since name alone isn't unique - def set_favorite(self, favorite_name_id: str): + def set_favorite(self, favorite_name_id: str) -> None: _LOGGER.debug(f"Setting favorite: {favorite_name_id}") fav_id = int(favorite_name_id.rsplit("-", 1)[1]) self._update({"current": {"srId": fav_id, "step": 1, "playing": "routine"}}) - def set_sound(self, sound_or_id_or_title: SoundContent | SimpleSoundContent | str | int | None, duration: int = 0, until="indefinite"): - """ - Set a sound by passing SoundContent item from self.sounds, id, or title. - """ + def set_sound( + self, + sound_or_id_or_title: SoundContent | SimpleSoundContent | str | int | None, + duration: int = 0, + until: IotSoundUntil = "indefinite", + ) -> None: + """Set a sound by passing SoundContent item from self.sounds, id, or title.""" if sound_or_id_or_title is None or sound_or_id_or_title == NO_SOUND_ID or sound_or_id_or_title == "none": self.turn_off() return @@ -189,7 +193,12 @@ def set_sound(self, sound_or_id_or_title: SoundContent | SimpleSoundContent | st # Assume it's a SoundContent or SimpleSoundContent object sound = sound_or_id_or_title - if not sound or not isinstance(sound, dict) or not sound.get('id') or (not sound.get('wavUrl') and not sound.get('mp3Url')): + if ( + not sound + or not isinstance(sound, dict) + or not sound.get("id") + or (not sound.get("wavUrl") and not sound.get("mp3Url")) + ): _LOGGER.error(f"Sound not found: {sound_or_id_or_title}") return @@ -201,7 +210,7 @@ def set_sound(self, sound_or_id_or_title: SoundContent | SimpleSoundContent | st "sound": { "id": sound["id"], "mute": False, - "url": sound.get("wavUrl") or sound.get("mp3Url"), + "url": sound.get("wavUrl") or cast(SoundContent, sound).get("mp3Url"), "duration": duration, "until": until, }, @@ -209,11 +218,11 @@ def set_sound(self, sound_or_id_or_title: SoundContent | SimpleSoundContent | st } ) - def turn_off(self): + def turn_off(self) -> None: _LOGGER.debug("Turning off sound") self._update({"current": {"srId": 0, "step": 0, "playing": "none"}}) - def turn_light_off(self): + def turn_light_off(self) -> None: _LOGGER.debug("Turning light off") # if favorite is playing then light can be turned off without turning off sound if self.current_playing == "routine": @@ -248,7 +257,7 @@ def turn_light_off(self): def set_color( self, red: int, green: int, blue: int, white: int = 0, brightness: int = 0 - ): + ) -> None: new_color_id = CUSTOM_COLOR_ID _LOGGER.debug( f"red: {red} green: {green} blue: {blue} brightness: {brightness} white: {white}" diff --git a/src/hatch_rest_api/shadow_client_subscriber.py b/src/hatch_rest_api/shadow_client_subscriber.py index e2cadd8..a785d4a 100644 --- a/src/hatch_rest_api/shadow_client_subscriber.py +++ b/src/hatch_rest_api/shadow_client_subscriber.py @@ -10,9 +10,8 @@ ShadowState, ) -from .types import SoundContent, SimpleSoundContent - from .callbacks import CallbacksMixin +from .types import JsonType, SimpleSoundContent, SoundContent _LOGGER = logging.getLogger(__name__) @@ -39,11 +38,13 @@ def __init__( self.shadow_client = shadow_client self.favorites = favorites self.sounds = sounds - self.sounds_by_id = {s['id']: s for s in sounds if s.get('id')} - self.sounds_by_name = {s['title']: s for s in sounds if s.get('title')} + self.sounds_by_id: dict[int, SoundContent | SimpleSoundContent] = {s["id"]: s for s in sounds if s.get("id")} + self.sounds_by_name: dict[str, SoundContent | SimpleSoundContent] = { + s["title"]: s for s in sounds if s.get("title") + } _LOGGER.debug(f"creating {self.__class__.__name__}: {device_name}") - def update_shadow_accepted(response: UpdateShadowResponse): + def update_shadow_accepted(response: UpdateShadowResponse) -> None: self._on_update_shadow_accepted(response) ( @@ -61,7 +62,7 @@ def update_shadow_accepted(response: UpdateShadowResponse): f"unsubscribe_topic_to_update_shadow_accepted: {unsubscribe_topic_to_update_shadow_accepted}" ) - def on_get_shadow_accepted(response: GetShadowResponse): + def on_get_shadow_accepted(response: GetShadowResponse) -> None: self._on_get_shadow_accepted(response) ( @@ -78,7 +79,7 @@ def on_get_shadow_accepted(response: GetShadowResponse): ) self.refresh() - def _on_update_shadow_accepted(self, response: UpdateShadowResponse): + def _on_update_shadow_accepted(self, response: UpdateShadowResponse) -> None: _LOGGER.debug(f"update {self.device_name}, RESPONSE: {response}") if response.version < self.document_version: _LOGGER.debug(f'ignoring update {self.device_name}, response version: {response.version} < document version: {self.document_version}') @@ -89,7 +90,7 @@ def _on_update_shadow_accepted(self, response: UpdateShadowResponse): self.document_version = response.version self._update_local_state(response.state.reported) - def _on_get_shadow_accepted(self, response: GetShadowResponse): + def _on_get_shadow_accepted(self, response: GetShadowResponse) -> None: _LOGGER.debug(f"get {self.device_name}, RESPONSE: {response}") if response.version < self.document_version: return @@ -101,7 +102,7 @@ def _on_get_shadow_accepted(self, response: GetShadowResponse): self.document_version = response.version self._update_local_state(response.state.reported) - def _update(self, desired_state): + def _update(self, desired_state: dict[str, JsonType]) -> None: _LOGGER.debug(f"updating: {desired_state}") request: UpdateShadowRequest = UpdateShadowRequest( thing_name=self.thing_name, @@ -113,7 +114,7 @@ def _update(self, desired_state): request, mqtt.QoS.AT_LEAST_ONCE ).result() - def refresh(self): + def refresh(self) -> None: _LOGGER.debug("Requesting current shadow state...") result = self.shadow_client.publish_get_shadow( request=iotshadow.GetShadowRequest( diff --git a/src/hatch_rest_api/stub.py b/src/hatch_rest_api/stub.py index b6206dc..c007d38 100644 --- a/src/hatch_rest_api/stub.py +++ b/src/hatch_rest_api/stub.py @@ -19,7 +19,7 @@ logger.addHandler(ch) -async def testing(): +async def testing() -> None: loop = asyncio.get_running_loop() email = input("Email: ") password = getpass() @@ -32,7 +32,7 @@ async def testing(): for iot_device in iot_devices: if "reading" in iot_device.device_name: - def output(): + def output() -> None: print(f"******-{iot_device}") iot_device.register_callback(output) diff --git a/src/hatch_rest_api/types.py b/src/hatch_rest_api/types.py index 7c7b931..15e86eb 100644 --- a/src/hatch_rest_api/types.py +++ b/src/hatch_rest_api/types.py @@ -1,5 +1,20 @@ -from typing import TypedDict, Any +from typing import Any, Literal, TypedDict +type JsonType = None | int | str | bool | list[JsonType] | dict[str, JsonType] + +type Product = Literal[ + "rest", + "riot", + "riotPlus", + "restPlus", + "restMini", + "restore", + "restoreIot", + "restoreV5", + "alexa", + "grow", + "answeredReader", +] class SimpleSoundContent(TypedDict): """ @@ -30,15 +45,15 @@ class SoundContent(TypedDict): tier: str extent: str alarmOnly: bool - author: "Any | None" - narrator: "Any | None" - duration: "Any | None" + author: Any | None + narrator: Any | None + duration: Any | None imageUrl: str - mp3Url: "str | None" - wavUrl: "str | None" + mp3Url: str | None + wavUrl: str | None color: Any - series: "list[Any]" - products: "list[str]" + series: list[Any] + products: list[str] libraryVersion: int libraryVersionString: str url: str @@ -46,18 +61,105 @@ class SoundContent(TypedDict): hatchId: int contentSource: str seriesSize: int - tagIds: "list" - altTagIds: "list" - mixIds: "list" - red: "Any | None" - green: "Any | None" - blue: "Any | None" - white: "Any | None" + tagIds: list + altTagIds: list + mixIds: list + red: Any | None + green: Any | None + blue: Any | None + white: Any | None sixCharacterColor: str rotateSeries: bool - contentfulId: "Any | None" + contentfulId: Any | None rotationType: str legacy: bool contentId: int contentful: bool personalized: bool + +type Gender = Literal['MALE', 'FEMALE'] + +class AwsIotCredentials(TypedDict): + AccessKeyId: str + Expiration: int + SecretKey: str + SessionToken: str + +class AwsIotCredentialsResponse(TypedDict): + Credentials: AwsIotCredentials + IdentityId: str + +class Baby(TypedDict): + id: int + createDate: str + updateDate: str + name: str + birthDate: str | None + dueDate: str | None + birthWeight: int | None + birthLength: int | None + gender: Gender + age: str + stage: int + lessThanTwoWeeksOld: bool + +class Member(TypedDict): + id: int + createDate: str + updateDate: str + active: bool + email: str + babies: list[Baby] + defaultUnitOfMeasure: Literal["imperial"] | str + appType: Literal["iOS"] | str + firstName: str + lastName: str | None + signupSource: Literal["rest"] | str + timezone: str + timeZoneAdjust: int + stage: int + +class LoginResponse(TypedDict): + payload: Member + sync: int + token: str + +class IotTokenResponse(TypedDict): + endpoint: str + identityId: str + region: str + cognitoPoolId: str + token: str + +class IotDeviceInfo(TypedDict): + id: int + createDate: str + updateDate: str + macAddress: str + owner: bool + name: str + hardwareVersion: str + product: Product + thingName: str + email: str + memberId: int + +type RestChargingStatus = Literal[0, 3, 5] + +class RestIotRoutine(TypedDict): + id: int + macAddress: str + name: str + type: Literal["favorite", "sleep", "wake", "flex"] + active: bool + enabled: bool + displayOrder: int + sleepScene: bool + followBySleepScene: bool + button0: bool # true if this routine is available on touch ring + startTime: None | str + endTime: None | str + daysOfWeek: Any | None + steps: list[Any] + +type IotSoundUntil = Literal["indefinite", "duration"] | str diff --git a/src/hatch_rest_api/util.py b/src/hatch_rest_api/util.py index 5a16262..a32467b 100644 --- a/src/hatch_rest_api/util.py +++ b/src/hatch_rest_api/util.py @@ -1,11 +1,14 @@ import logging +from collections.abc import Callable +from typing import Any, overload from .const import SENSITIVE_FIELD_NAMES, MAX_IOT_VALUE +from .types import JsonType _LOGGER = logging.getLogger(__name__) -def clean_dictionary_for_logging(dictionary: dict[str, any]) -> dict[str, any]: +def clean_dictionary_for_logging(dictionary: dict[str, Any]) -> dict[str, Any]: mutable_dictionary = dictionary.copy() for key in dictionary: if key.lower() in SENSITIVE_FIELD_NAMES: @@ -26,24 +29,29 @@ def clean_dictionary_for_logging(dictionary: dict[str, any]) -> dict[str, any]: return mutable_dictionary -def convert_from_percentage(percentage: int): +def convert_from_percentage(percentage: float) -> int: return int(round(percentage / 100.0 * MAX_IOT_VALUE)) & 0xFFFF -def convert_to_percentage(value: int): +def convert_to_percentage(value: int) -> int: return round((value & 0xFFFF) / (1.0 * MAX_IOT_VALUE) * 100) -def convert_from_hex(percentage: int): +def convert_from_hex(percentage: int) -> int: return int(round(percentage / 255.0 * MAX_IOT_VALUE)) & 0xFFFF -def convert_to_hex(value: int): +def convert_to_hex(value: int) -> int: return round((value & 0xFFFF) / (1.0 * MAX_IOT_VALUE) * 255) -def safely_get_json_value(json, key, callable_to_cast=None): - value = json +@overload +def safely_get_json_value[T](json: dict[str, JsonType], key: str, callable_to_cast: Callable[..., T]) -> T | None: ... +@overload +def safely_get_json_value(json: dict[str, JsonType], key: str, callable_to_cast: None = None) -> JsonType: ... + +def safely_get_json_value[T](json: dict[str, JsonType], key: str, callable_to_cast: Callable[..., T] | None = None) -> T | JsonType | None: + value: JsonType = json for x in key.split("."): if value is not None: try: @@ -54,5 +62,5 @@ def safely_get_json_value(json, key, callable_to_cast=None): except (TypeError, KeyError, ValueError): value = None if callable_to_cast is not None and value is not None: - value = callable_to_cast(value) + return callable_to_cast(value) return value diff --git a/src/hatch_rest_api/util_bootstrap.py b/src/hatch_rest_api/util_bootstrap.py index 6f10e8c..52f5bf8 100644 --- a/src/hatch_rest_api/util_bootstrap.py +++ b/src/hatch_rest_api/util_bootstrap.py @@ -2,15 +2,18 @@ import logging from functools import partial from re import IGNORECASE, sub +from typing import Protocol from uuid import uuid4 from aiohttp import ClientSession from awscrt import io from awscrt.auth import AwsCredentialsProvider +from awscrt.exceptions import AwsCrtError +from awscrt.mqtt import Connection, ConnectReturnCode from awsiot.iotshadow import IotShadowClient from awsiot.mqtt_connection_builder import websockets_with_default_aws_signing -from . import BaseError +from . import BaseError, RestDevice from .aws_http import AwsHttp from .const import NO_SOUND_ID from .contentful import Contentful @@ -21,18 +24,32 @@ from .rest_plus import RestPlus from .restore_iot import RestoreIot from .restore_v5 import RestoreV5 -from .types import SimpleSoundContent +from .types import IotDeviceInfo, RestIotRoutine, SimpleSoundContent _LOGGER = logging.getLogger(__name__) +class ConnectionInterruptedCallback(Protocol): + def __call__(self, *, connection: Connection, error: AwsCrtError) -> None: ... + + +class ConnectionResumedCallback(Protocol): + def __call__( + self, + *, + connection: Connection, + return_code: ConnectReturnCode, + session_present: bool, + ) -> None: ... + + async def get_rest_devices( email: str, password: str, - client_session: ClientSession = None, - on_connection_interrupted=None, - on_connection_resumed=None, -): + client_session: ClientSession | None = None, + on_connection_interrupted: ConnectionInterruptedCallback | None = None, + on_connection_resumed: ConnectionResumedCallback | None = None, +) -> tuple[Hatch, Connection, list[RestDevice], int]: loop = asyncio.get_running_loop() if _LOGGER.isEnabledFor(logging.DEBUG): await loop.run_in_executor( @@ -91,7 +108,7 @@ async def get_rest_devices( shadow_client = IotShadowClient(mqtt_connection) - def create_rest_devices(iot_device): + def create_rest_devices(iot_device: IotDeviceInfo) -> RestDevice: mac_address = iot_device["macAddress"] if mac_address in favorites_map: favorites = list(favorites_map[mac_address]) @@ -154,7 +171,9 @@ def create_rest_devices(iot_device): ) -async def _get_favorites_for_all_v2_devices(api, token, iot_devices): +async def _get_favorites_for_all_v2_devices( + api: Hatch, token: str, iot_devices: list[IotDeviceInfo] +) -> dict[str, list[RestIotRoutine]]: mac_to_favorite = {} for device in iot_devices: if device["product"] in ["riot", "riotPlus", "restoreV5"]: @@ -165,7 +184,9 @@ async def _get_favorites_for_all_v2_devices(api, token, iot_devices): return mac_to_favorite -async def _get_routines_for_all_v2_devices(api, token, iot_devices): +async def _get_routines_for_all_v2_devices( + api: Hatch, token: str, iot_devices: list[IotDeviceInfo] +) -> dict[str, list[RestIotRoutine]]: mac_to_routines = {} for device in iot_devices: if device["product"] in ["riot", "restoreIot", "restoreV5"]: @@ -177,7 +198,7 @@ async def _get_routines_for_all_v2_devices(api, token, iot_devices): async def _get_sound_content_for_all_v2_devices( - api: Hatch, token: str, contentful: Contentful, iot_devices + api: Hatch, token: str, contentful: Contentful, iot_devices: list[IotDeviceInfo] ) -> dict[str, list[SimpleSoundContent]]: mac_to_sounds = {} for device in iot_devices: @@ -187,7 +208,9 @@ async def _get_sound_content_for_all_v2_devices( content = await api.content( auth_token=token, product="riot", content=["sound"] ) - sounds = [s for s in content["contentItems"] if s["id"] != NO_SOUND_ID] + sounds: list[SimpleSoundContent] = [ + s for s in content["contentItems"] if s["id"] != NO_SOUND_ID + ] except RateError as e: _LOGGER.warning( f"Rate limit error when fetching sounds for {mac}: {str(e)}" diff --git a/src/hatch_rest_api/util_http.py b/src/hatch_rest_api/util_http.py index 786e49f..be0cda9 100644 --- a/src/hatch_rest_api/util_http.py +++ b/src/hatch_rest_api/util_http.py @@ -1,18 +1,23 @@ +from collections.abc import Awaitable, Callable import logging +from typing import cast +from aiohttp import ClientResponse + +from .types import JsonType from .util import clean_dictionary_for_logging _LOGGER = logging.getLogger(__name__) -def request_with_logging(func): - async def request_with_logging_wrapper(*args, **kwargs): +def request_with_logging[**P, T: ClientResponse](func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]: + async def request_with_logging_wrapper(*args: P.args, **kwargs: P.kwargs) -> T: url = kwargs["url"] request_message = f"sending {url} request" headers = kwargs.get("headers") if headers is not None: request_message = request_message + f"headers: {headers}" - json_body = kwargs.get("json_body") + json_body = cast(dict[str, JsonType], kwargs.get("json_body")) if json_body is not None: request_message = ( request_message From 9df31b534ed5076ac8fb4bdea8f1b22b6fbab00a Mon Sep 17 00:00:00 2001 From: Bao Trinh Date: Wed, 6 Aug 2025 23:00:32 -0500 Subject: [PATCH 2/2] fix: type hints fixes with runtime changes --- src/hatch_rest_api/restore_v5.py | 15 +-- .../shadow_client_subscriber.py | 6 +- src/hatch_rest_api/util_bootstrap.py | 113 +++++++++++------- 3 files changed, 81 insertions(+), 53 deletions(-) diff --git a/src/hatch_rest_api/restore_v5.py b/src/hatch_rest_api/restore_v5.py index 4faac10..8e49600 100644 --- a/src/hatch_rest_api/restore_v5.py +++ b/src/hatch_rest_api/restore_v5.py @@ -185,13 +185,14 @@ def set_sound( self.turn_off() return - if isinstance(sound_or_id_or_title, int): - sound = self.sounds_by_id.get(sound_or_id_or_title) - elif isinstance(sound_or_id_or_title, str): - sound = self.sounds_by_name.get(sound_or_id_or_title) - else: - # Assume it's a SoundContent or SimpleSoundContent object - sound = sound_or_id_or_title + match sound_or_id_or_title: + case int(): + sound = self.sounds_by_id.get(sound_or_id_or_title) + case str(): + sound = self.sounds_by_name.get(sound_or_id_or_title) + case dict(): + # Assume it's a SoundContent or SimpleSoundContent object + sound = sound_or_id_or_title if ( not sound diff --git a/src/hatch_rest_api/shadow_client_subscriber.py b/src/hatch_rest_api/shadow_client_subscriber.py index a785d4a..1da24a1 100644 --- a/src/hatch_rest_api/shadow_client_subscriber.py +++ b/src/hatch_rest_api/shadow_client_subscriber.py @@ -1,3 +1,4 @@ +import abc import logging from awscrt import mqtt @@ -16,7 +17,7 @@ _LOGGER = logging.getLogger(__name__) -class ShadowClientSubscriberMixin(CallbacksMixin): +class ShadowClientSubscriberMixin(CallbacksMixin, abc.ABC): document_version: int = -1 def __init__( @@ -79,6 +80,9 @@ def on_get_shadow_accepted(response: GetShadowResponse) -> None: ) self.refresh() + @abc.abstractmethod + def _update_local_state(self, state: dict[str, JsonType]) -> None: ... + def _on_update_shadow_accepted(self, response: UpdateShadowResponse) -> None: _LOGGER.debug(f"update {self.device_name}, RESPONSE: {response}") if response.version < self.document_version: diff --git a/src/hatch_rest_api/util_bootstrap.py b/src/hatch_rest_api/util_bootstrap.py index 52f5bf8..60d8866 100644 --- a/src/hatch_rest_api/util_bootstrap.py +++ b/src/hatch_rest_api/util_bootstrap.py @@ -2,7 +2,7 @@ import logging from functools import partial from re import IGNORECASE, sub -from typing import Protocol +from typing import Protocol, TypedDict, cast from uuid import uuid4 from aiohttp import ClientSession @@ -120,47 +120,51 @@ def create_rest_devices(iot_device: IotDeviceInfo) -> RestDevice: else: _LOGGER.debug(f"Iot device {iot_device} has no routines") routines = [] - if iot_device["product"] == "restPlus": - return RestPlus( - device_name=iot_device["name"], - thing_name=iot_device["thingName"], - mac=mac_address, - shadow_client=shadow_client, - ) - elif iot_device["product"] in ["riot", "riotPlus"]: - return RestIot( - device_name=iot_device["name"], - thing_name=iot_device["thingName"], - mac=mac_address, - shadow_client=shadow_client, - favorites=favorites, - sounds=sounds_map[mac_address], - ) - elif iot_device["product"] == "restoreIot": - return RestoreIot( - device_name=iot_device["name"], - thing_name=iot_device["thingName"], - mac=mac_address, - shadow_client=shadow_client, - favorites=routines + favorites, - sounds=sounds_map[mac_address], - ) - elif iot_device["product"] == "restoreV5": - return RestoreV5( - device_name=iot_device["name"], - thing_name=iot_device["thingName"], - mac=mac_address, - shadow_client=shadow_client, - favorites=routines + favorites, - sounds=sounds_map[mac_address], - ) - else: - return RestMini( - device_name=iot_device["name"], - thing_name=iot_device["thingName"], - mac=mac_address, - shadow_client=shadow_client, - ) + + match iot_device["product"]: + case "restPlus": + return RestPlus( + device_name=iot_device["name"], + thing_name=iot_device["thingName"], + mac=mac_address, + shadow_client=shadow_client, + ) + case "riot" | "riotPlus": + return RestIot( + device_name=iot_device["name"], + thing_name=iot_device["thingName"], + mac=mac_address, + shadow_client=shadow_client, + favorites=favorites, + sounds=sounds_map[mac_address], + ) + case "restoreIot": + return RestoreIot( + device_name=iot_device["name"], + thing_name=iot_device["thingName"], + mac=mac_address, + shadow_client=shadow_client, + favorites=routines + favorites, + sounds=sounds_map[mac_address], + ) + case "restoreV5": + return RestoreV5( + device_name=iot_device["name"], + thing_name=iot_device["thingName"], + mac=mac_address, + shadow_client=shadow_client, + favorites=routines + favorites, + sounds=sounds_map[mac_address], + ) + case "restMini": + return RestMini( + device_name=iot_device["name"], + thing_name=iot_device["thingName"], + mac=mac_address, + shadow_client=shadow_client, + ) + case _: + raise BaseError(f"Unsupported device type: {iot_device['product']}") rest_devices = map(create_rest_devices, iot_devices) return ( @@ -218,7 +222,7 @@ async def _get_sound_content_for_all_v2_devices( sounds = [] elif device["product"] == "restoreV5": try: - content = await contentful.graphql_query( + gql_content = await contentful.graphql_query( auth_token=token, query=""" query GetSounds($product: String!) { @@ -253,9 +257,28 @@ async def _get_sound_content_for_all_v2_devices( """, product=device["product"], ) + + class GqlSoundFile(TypedDict): + url: str + + class GqlSoundContent(TypedDict): + id: int + title: str + wavFile: GqlSoundFile + + class GqlSoundCollection(TypedDict): + total: int + limit: int + items: list[GqlSoundContent] + + class GqlSoundCollectionResponse(TypedDict): + soundCollection: GqlSoundCollection + sounds = [ - {**s, "wavUrl": s["wavFile"]["url"]} - for s in content["soundCollection"]["items"] + {"id": s["id"], "title": s["title"], "wavUrl": s["wavFile"]["url"]} + for s in cast(GqlSoundCollectionResponse, gql_content)[ + "soundCollection" + ]["items"] if s["id"] != NO_SOUND_ID ] except RateError as e: