diff --git a/pylint/plugins/hass_enforce_type_hints.py b/pylint/plugins/hass_enforce_type_hints.py index 3d86767df43..a6c4a968aa7 100644 --- a/pylint/plugins/hass_enforce_type_hints.py +++ b/pylint/plugins/hass_enforce_type_hints.py @@ -577,6 +577,20 @@ _ENTITY_MATCH: list[TypeHintMatch] = [ has_async_counterpart=True, ), ] +_RESTORE_ENTITY_MATCH: list[TypeHintMatch] = [ + TypeHintMatch( + function_name="async_get_last_state", + return_type=["State", None], + ), + TypeHintMatch( + function_name="async_get_last_extra_data", + return_type=["ExtraStoredData", None], + ), + TypeHintMatch( + function_name="extra_restore_state_data", + return_type=["ExtraStoredData", None], + ), +] _TOGGLE_ENTITY_MATCH: list[TypeHintMatch] = [ TypeHintMatch( function_name="is_on", @@ -1786,6 +1800,10 @@ _INHERITANCE_MATCH: dict[str, list[ClassTypeHintMatch]] = { base_class="Entity", matches=_ENTITY_MATCH, ), + ClassTypeHintMatch( + base_class="RestoreEntity", + matches=_RESTORE_ENTITY_MATCH, + ), ClassTypeHintMatch( base_class="NumberEntity", matches=[ @@ -1829,6 +1847,19 @@ _INHERITANCE_MATCH: dict[str, list[ClassTypeHintMatch]] = { ), ], ), + ClassTypeHintMatch( + base_class="RestoreNumber", + matches=[ + TypeHintMatch( + function_name="extra_restore_state_data", + return_type="NumberExtraStoredData", + ), + TypeHintMatch( + function_name="async_get_last_number_data", + return_type=["NumberExtraStoredData", None], + ), + ], + ), ], "remote": [ ClassTypeHintMatch(