Add platform Entity classes to pylint plugin (#125737)

* Add platform Entity classes to pylint plugin

* Fix violations

* Fix violations

* More

* Allow component package with same name as a platform

* One more

homeassistant/components/hue/v1/binary_sensor.py

@@ -25,6 +25,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     )
 
 
+# pylint: disable-next=hass-enforce-class-module
 class HuePresence(GenericZLLSensor, BinarySensorEntity):
     """The presence sensor entity for a Hue motion sensor device."""
 

homeassistant/components/hue/v1/light.py

@@ -305,6 +305,7 @@ def hass_to_hue_brightness(value):
     return max(1, round((value / 255) * 254))
 
 
+# pylint: disable-next=hass-enforce-class-module
 class HueLight(CoordinatorEntity, LightEntity):
     """Representation of a Hue light."""
 

homeassistant/components/hue/v1/sensor.py

@@ -32,10 +32,12 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     await bridge.sensor_manager.async_register_component("sensor", async_add_entities)
 
 
+# pylint: disable-next=hass-enforce-class-module
 class GenericHueGaugeSensorEntity(GenericZLLSensor, SensorEntity):
     """Parent class for all 'gauge' Hue device sensors."""
 
 
+# pylint: disable-next=hass-enforce-class-module
 class HueLightLevel(GenericHueGaugeSensorEntity):
     """The light level sensor entity for a Hue motion sensor device."""
 
@@ -71,6 +73,7 @@ class HueLightLevel(GenericHueGaugeSensorEntity):
         return attributes
 
 
+# pylint: disable-next=hass-enforce-class-module
 class HueTemperature(GenericHueGaugeSensorEntity):
     """The temperature sensor entity for a Hue motion sensor device."""
 
@@ -87,6 +90,7 @@ class HueTemperature(GenericHueGaugeSensorEntity):
         return self.sensor.temperature / 100
 
 
+# pylint: disable-next=hass-enforce-class-module
 class HueBattery(GenericHueSensor, SensorEntity):
     """Battery class for when a batt-powered device is only represented as an event."""
 

homeassistant/components/hue/v2/binary_sensor.py

@@ -82,6 +82,7 @@ async def async_setup_entry(
     register_items(api.sensors.tamper, HueTamperSensor)
 
 
+# pylint: disable-next=hass-enforce-class-module
 class HueMotionSensor(HueBaseEntity, BinarySensorEntity):
     """Representation of a Hue Motion sensor."""
 
@@ -103,6 +104,7 @@ class HueMotionSensor(HueBaseEntity, BinarySensorEntity):
         return self.resource.motion.value
 
 
+# pylint: disable-next=hass-enforce-class-module
 class HueEntertainmentActiveSensor(HueBaseEntity, BinarySensorEntity):
     """Representation of a Hue Entertainment Configuration as binary sensor."""
 
@@ -126,6 +128,7 @@ class HueEntertainmentActiveSensor(HueBaseEntity, BinarySensorEntity):
         return self.resource.metadata.name
 
 
+# pylint: disable-next=hass-enforce-class-module
 class HueContactSensor(HueBaseEntity, BinarySensorEntity):
     """Representation of a Hue Contact sensor."""
 
@@ -147,6 +150,7 @@ class HueContactSensor(HueBaseEntity, BinarySensorEntity):
         return self.resource.contact_report.state != ContactState.CONTACT
 
 
+# pylint: disable-next=hass-enforce-class-module
 class HueTamperSensor(HueBaseEntity, BinarySensorEntity):
     """Representation of a Hue Tamper sensor."""
 

homeassistant/components/hue/v2/group.py

@@ -76,6 +76,7 @@ async def async_setup_entry(
     )
 
 
+# pylint: disable-next=hass-enforce-class-module
 class GroupedHueLight(HueBaseEntity, LightEntity):
     """Representation of a Grouped Hue light."""
 

homeassistant/components/hue/v2/light.py

@@ -68,6 +68,7 @@ async def async_setup_entry(
     )
 
 
+# pylint: disable-next=hass-enforce-class-module
 class HueLight(HueBaseEntity, LightEntity):
     """Representation of a Hue light."""
 

homeassistant/components/hue/v2/sensor.py

@@ -79,6 +79,7 @@ async def async_setup_entry(
     register_items(ctrl_base.zigbee_connectivity, HueZigbeeConnectivitySensor)
 
 
+# pylint: disable-next=hass-enforce-class-module
 class HueSensorBase(HueBaseEntity, SensorEntity):
     """Representation of a Hue sensor."""
 
@@ -94,6 +95,7 @@ class HueSensorBase(HueBaseEntity, SensorEntity):
         self.controller = controller
 
 
+# pylint: disable-next=hass-enforce-class-module
 class HueTemperatureSensor(HueSensorBase):
     """Representation of a Hue Temperature sensor."""
 
@@ -111,6 +113,7 @@ class HueTemperatureSensor(HueSensorBase):
         return round(self.resource.temperature.value, 1)
 
 
+# pylint: disable-next=hass-enforce-class-module
 class HueLightLevelSensor(HueSensorBase):
     """Representation of a Hue LightLevel (illuminance) sensor."""
 
@@ -139,6 +142,7 @@ class HueLightLevelSensor(HueSensorBase):
         }
 
 
+# pylint: disable-next=hass-enforce-class-module
 class HueBatterySensor(HueSensorBase):
     """Representation of a Hue Battery sensor."""
 
@@ -164,6 +168,7 @@ class HueBatterySensor(HueSensorBase):
         return {"battery_state": self.resource.power_state.battery_state.value}
 
 
+# pylint: disable-next=hass-enforce-class-module
 class HueZigbeeConnectivitySensor(HueSensorBase):
     """Representation of a Hue ZigbeeConnectivity sensor."""
 

homeassistant/components/input_button/__init__.py

@@ -128,6 +128,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
+# pylint: disable-next=hass-enforce-class-module
 class InputButton(collection.CollectionEntity, ButtonEntity, RestoreEntity):
     """Representation of a button."""
 

homeassistant/components/input_select/__init__.py

@@ -246,6 +246,7 @@ class InputSelectStorageCollection(collection.DictStorageCollection):
         return {CONF_ID: item[CONF_ID]} | update_data
 
 
+# pylint: disable-next=hass-enforce-class-module
 class InputSelect(collection.CollectionEntity, SelectEntity, RestoreEntity):
     """Representation of a select input."""
 

pylint/plugins/hass_enforce_class_module.py

@@ -3,49 +3,66 @@
 from __future__ import annotations
 
 from ast import ClassDef
-from dataclasses import dataclass
 
 from astroid import nodes
 from pylint.checkers import BaseChecker
 from pylint.lint import PyLinter
 
-
+_MODULES: dict[str, set[str]] = {
-@dataclass
+    "air_quality": {"AirQualityEntity"},
-class ClassModuleMatch:
+    "alarm_control_panel": {
-    """Class for pattern matching."""
+        "AlarmControlPanelEntity",
-
+        "AlarmControlPanelEntityDescription",
-    expected_module: str
+    },
-    base_class: str
+    "assist_satellite": {"AssistSatelliteEntity", "AssistSatelliteEntityDescription"},
-
+    "binary_sensor": {"BinarySensorEntity", "BinarySensorEntityDescription"},
-
+    "button": {"ButtonEntity", "ButtonEntityDescription"},
-_MODULES = [
+    "calendar": {"CalendarEntity"},
-    ClassModuleMatch("alarm_control_panel", "AlarmControlPanelEntityDescription"),
+    "camera": {"CameraEntity", "CameraEntityDescription"},
-    ClassModuleMatch("assist_satellite", "AssistSatelliteEntityDescription"),
+    "climate": {"ClimateEntity", "ClimateEntityDescription"},
-    ClassModuleMatch("binary_sensor", "BinarySensorEntityDescription"),
+    "coordinator": {"DataUpdateCoordinator"},
-    ClassModuleMatch("button", "ButtonEntityDescription"),
+    "conversation": {"ConversationEntity"},
-    ClassModuleMatch("camera", "CameraEntityDescription"),
+    "cover": {"CoverEntity", "CoverEntityDescription"},
-    ClassModuleMatch("climate", "ClimateEntityDescription"),
+    "date": {"DateEntity", "DateEntityDescription"},
-    ClassModuleMatch("coordinator", "DataUpdateCoordinator"),
+    "datetime": {"DateTimeEntity", "DateTimeEntityDescription"},
-    ClassModuleMatch("cover", "CoverEntityDescription"),
+    "device_tracker": {"DeviceTrackerEntity"},
-    ClassModuleMatch("date", "DateEntityDescription"),
+    "event": {"EventEntity", "EventEntityDescription"},
-    ClassModuleMatch("datetime", "DateTimeEntityDescription"),
+    "fan": {"FanEntity", "FanEntityDescription"},
-    ClassModuleMatch("event", "EventEntityDescription"),
+    "geo_location": {"GeolocationEvent"},
-    ClassModuleMatch("image", "ImageEntityDescription"),
+    "humidifier": {"HumidifierEntity", "HumidifierEntityDescription"},
-    ClassModuleMatch("image_processing", "ImageProcessingEntityDescription"),
+    "image": {"ImageEntity", "ImageEntityDescription"},
-    ClassModuleMatch("lawn_mower", "LawnMowerEntityDescription"),
+    "image_processing": {
-    ClassModuleMatch("lock", "LockEntityDescription"),
+        "ImageProcessingEntity",
-    ClassModuleMatch("media_player", "MediaPlayerEntityDescription"),
+        "ImageProcessingFaceEntity",
-    ClassModuleMatch("notify", "NotifyEntityDescription"),
+        "ImageProcessingEntityDescription",
-    ClassModuleMatch("number", "NumberEntityDescription"),
+    },
-    ClassModuleMatch("select", "SelectEntityDescription"),
+    "lawn_mower": {"LawnMowerEntity", "LawnMowerEntityDescription"},
-    ClassModuleMatch("sensor", "SensorEntityDescription"),
+    "light": {"LightEntity", "LightEntityDescription"},
-    ClassModuleMatch("text", "TextEntityDescription"),
+    "lock": {"LockEntity", "LockEntityDescription"},
-    ClassModuleMatch("time", "TimeEntityDescription"),
+    "media_player": {"MediaPlayerEntity", "MediaPlayerEntityDescription"},
-    ClassModuleMatch("update", "UpdateEntityDescription"),
+    "notify": {"NotifyEntity", "NotifyEntityDescription"},
-    ClassModuleMatch("vacuum", "VacuumEntityDescription"),
+    "number": {"NumberEntity", "NumberEntityDescription", "RestoreNumber"},
-    ClassModuleMatch("water_heater", "WaterHeaterEntityDescription"),
+    "remote": {"RemoteEntity", "RemoteEntityDescription"},
-    ClassModuleMatch("weather", "WeatherEntityDescription"),
+    "select": {"SelectEntity", "SelectEntityDescription"},
-]
+    "sensor": {"RestoreSensor", "SensorEntity", "SensorEntityDescription"},
+    "siren": {"SirenEntity", "SirenEntityDescription"},
+    "stt": {"SpeechToTextEntity"},
+    "switch": {"SwitchEntity", "SwitchEntityDescription"},
+    "text": {"TextEntity", "TextEntityDescription"},
+    "time": {"TimeEntity", "TimeEntityDescription"},
+    "todo": {"TodoListEntity"},
+    "tts": {"TextToSpeechEntity"},
+    "update": {"UpdateEntityDescription"},
+    "vacuum": {"VacuumEntity", "VacuumEntityDescription"},
+    "wake_word": {"WakeWordDetectionEntity"},
+    "water_heater": {"WaterHeaterEntity"},
+    "weather": {
+        "CoordinatorWeatherEntity",
+        "SingleCoordinatorWeatherEntity",
+        "WeatherEntity",
+        "WeatherEntityDescription",
+    },
+}
 
 
 class HassEnforceClassModule(BaseChecker):
@@ -69,24 +86,24 @@ class HassEnforceClassModule(BaseChecker):
         if not root_name.startswith("homeassistant.components."):
             return
         parts = root_name.split(".")
+        current_integration = parts[2]
         current_module = parts[3] if len(parts) > 3 else ""
 
         ancestors: list[ClassDef] | None = None
 
-        for match in _MODULES:
+        for expected_module, classes in _MODULES.items():
-            # Allow module.py and module/sub_module.py
+            if expected_module in (current_module, current_integration):
-            if current_module == match.expected_module:
                 continue
 
             if ancestors is None:
                 ancestors = list(node.ancestors())  # cache result for other modules
 
             for ancestor in ancestors:
-                if ancestor.name == match.base_class:
+                if ancestor.name in classes:
                     self.add_message(
                         "hass-enforce-class-module",
                         node=node,
-                        args=(match.base_class, match.expected_module),
+                        args=(ancestor.name, expected_module),
                     )
                     return
 

tests/pylint/test_enforce_class_module.py

@@ -63,6 +63,36 @@ def test_enforce_class_module_good(
         walker.walk(root_node)
 
 
+@pytest.mark.parametrize(
+    "path",
+    [
+        "homeassistant.components.sensor",
+        "homeassistant.components.sensor.entity",
+        "homeassistant.components.pylint_test.sensor",
+        "homeassistant.components.pylint_test.sensor.entity",
+    ],
+)
+def test_enforce_class_platform_good(
+    linter: UnittestLinter,
+    enforce_class_module_checker: BaseChecker,
+    path: str,
+) -> None:
+    """Good test cases."""
+    code = """
+    class SensorEntity:
+        pass
+
+    class CustomSensorEntity(SensorEntity):
+        pass
+    """
+    root_node = astroid.parse(code, path)
+    walker = ASTWalker(linter)
+    walker.add_checker(enforce_class_module_checker)
+
+    with assert_no_messages(linter):
+        walker.walk(root_node)
+
+
 @pytest.mark.parametrize(
     "path",
     [