From 27b5ebb9d3a1912f55ec365b805f6e15f8bfe1a8 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Wed, 17 Aug 2022 18:43:28 +0200
Subject: [PATCH] Add RestoreSensor to sensor checks in pylint (#76916)

---
 pylint/plugins/hass_enforce_type_hints.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/pylint/plugins/hass_enforce_type_hints.py b/pylint/plugins/hass_enforce_type_hints.py
index 4e81b7b77af..7eb31eba523 100644
--- a/pylint/plugins/hass_enforce_type_hints.py
+++ b/pylint/plugins/hass_enforce_type_hints.py
@@ -1949,6 +1949,10 @@ _INHERITANCE_MATCH: dict[str, list[ClassTypeHintMatch]] = {
             base_class="Entity",
             matches=_ENTITY_MATCH,
         ),
+        ClassTypeHintMatch(
+            base_class="RestoreEntity",
+            matches=_RESTORE_ENTITY_MATCH,
+        ),
         ClassTypeHintMatch(
             base_class="SensorEntity",
             matches=[
@@ -1983,6 +1987,19 @@ _INHERITANCE_MATCH: dict[str, list[ClassTypeHintMatch]] = {
                 ),
             ],
         ),
+        ClassTypeHintMatch(
+            base_class="RestoreSensor",
+            matches=[
+                TypeHintMatch(
+                    function_name="extra_restore_state_data",
+                    return_type="SensorExtraStoredData",
+                ),
+                TypeHintMatch(
+                    function_name="async_get_last_sensor_data",
+                    return_type=["SensorExtraStoredData", None],
+                ),
+            ],
+        ),
     ],
     "siren": [
         ClassTypeHintMatch(