Add EntityDescription classes to pylint plugin (#125596)

* Add EntityDescription classes to pylint plugin

* Ignore existing violations

* Adjust

homeassistant/components/dsmr_reader/definitions.py

@@ -40,6 +40,7 @@ def tariff_transform(value: str) -> str:
 
 
 @dataclass(frozen=True)
+# pylint: disable-next=hass-enforce-class-module
 class DSMRReaderSensorEntityDescription(SensorEntityDescription):
     """Sensor entity description for DSMR Reader."""
 

homeassistant/components/growatt_server/sensor_types/sensor_entity_description.py

@@ -15,6 +15,7 @@ class GrowattRequiredKeysMixin:
 
 
 @dataclass(frozen=True)
+# pylint: disable-next=hass-enforce-class-module
 class GrowattSensorEntityDescription(SensorEntityDescription, GrowattRequiredKeysMixin):
     """Describes Growatt sensor entity."""
 

homeassistant/components/repetier/__init__.py

@@ -133,6 +133,7 @@ class RepetierRequiredKeysMixin:
 
 
 @dataclass(frozen=True)
+# pylint: disable-next=hass-enforce-class-module
 class RepetierSensorEntityDescription(
     SensorEntityDescription, RepetierRequiredKeysMixin
 ):

homeassistant/components/sunweg/sensor_types/sensor_entity_description.py

@@ -15,6 +15,7 @@ class SunWEGRequiredKeysMixin:
 
 
 @dataclass(frozen=True)
+# pylint: disable-next=hass-enforce-class-module
 class SunWEGSensorEntityDescription(SensorEntityDescription, SunWEGRequiredKeysMixin):
     """Describes SunWEG sensor entity."""
 

pylint/plugins/hass_enforce_class_module.py

@@ -1,37 +1,90 @@
-"""Plugin for checking if coordinator is in its own module."""
+"""Plugin for checking if class is in correct module."""
 
 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
 
 
+@dataclass
+class ClassModuleMatch:
+    """Class for pattern matching."""
+
+    expected_module: str
+    base_class: str
+
+
+_MODULES = [
+    ClassModuleMatch("alarm_control_panel", "AlarmControlPanelEntityDescription"),
+    ClassModuleMatch("assist_satellite", "AssistSatelliteEntityDescription"),
+    ClassModuleMatch("binary_sensor", "BinarySensorEntityDescription"),
+    ClassModuleMatch("button", "ButtonEntityDescription"),
+    ClassModuleMatch("camera", "CameraEntityDescription"),
+    ClassModuleMatch("climate", "ClimateEntityDescription"),
+    ClassModuleMatch("coordinator", "DataUpdateCoordinator"),
+    ClassModuleMatch("cover", "CoverEntityDescription"),
+    ClassModuleMatch("date", "DateEntityDescription"),
+    ClassModuleMatch("datetime", "DateTimeEntityDescription"),
+    ClassModuleMatch("event", "EventEntityDescription"),
+    ClassModuleMatch("image", "ImageEntityDescription"),
+    ClassModuleMatch("image_processing", "ImageProcessingEntityDescription"),
+    ClassModuleMatch("lawn_mower", "LawnMowerEntityDescription"),
+    ClassModuleMatch("lock", "LockEntityDescription"),
+    ClassModuleMatch("media_player", "MediaPlayerEntityDescription"),
+    ClassModuleMatch("notify", "NotifyEntityDescription"),
+    ClassModuleMatch("number", "NumberEntityDescription"),
+    ClassModuleMatch("select", "SelectEntityDescription"),
+    ClassModuleMatch("sensor", "SensorEntityDescription"),
+    ClassModuleMatch("text", "TextEntityDescription"),
+    ClassModuleMatch("time", "TimeEntityDescription"),
+    ClassModuleMatch("update", "UpdateEntityDescription"),
+    ClassModuleMatch("vacuum", "VacuumEntityDescription"),
+    ClassModuleMatch("water_heater", "WaterHeaterEntityDescription"),
+    ClassModuleMatch("weather", "WeatherEntityDescription"),
+]
+
+
 class HassEnforceClassModule(BaseChecker):
-    """Checker for coordinators own module."""
+    """Checker for class in correct module."""
 
     name = "hass_enforce_class_module"
     priority = -1
     msgs = {
         "C7461": (
-            "Derived data update coordinator is recommended to be placed in the 'coordinator' module",
+            "Derived %s is recommended to be placed in the '%s' module",
             "hass-enforce-class-module",
-            "Used when derived data update coordinator should be placed in its own module.",
+            "Used when derived class should be placed in its own module.",
         ),
     }
 
     def visit_classdef(self, node: nodes.ClassDef) -> None:
-        """Check if derived data update coordinator is placed in its own module."""
+        """Check if derived class is placed in its own module."""
         root_name = node.root().name
 
-        # we only want to check component update coordinators
+        # we only want to check components
-        if not root_name.startswith("homeassistant.components"):
+        if not root_name.startswith("homeassistant.components."):
             return
 
-        is_coordinator_module = root_name.endswith(".coordinator")
+        ancestors: list[ClassDef] | None = None
-        for ancestor in node.ancestors():
+
-            if ancestor.name == "DataUpdateCoordinator" and not is_coordinator_module:
+        for match in _MODULES:
-                self.add_message("hass-enforce-class-module", node=node)
+            if root_name.endswith(f".{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:
+                    self.add_message(
+                        "hass-enforce-class-module",
+                        node=node,
+                        args=(match.base_class, match.expected_module),
+                    )
                     return
 
 

tests/pylint/conftest.py

@@ -113,13 +113,11 @@ def hass_enforce_class_module_fixture() -> ModuleType:
     )
 
 
-@pytest.fixture(name="enforce_coordinator_module_checker")
+@pytest.fixture(name="enforce_class_module_checker")
-def enforce_coordinator_module_fixture(
+def enforce_class_module_fixture(hass_enforce_class_module, linter) -> BaseChecker:
-    hass_enforce_class_module, linter
-) -> BaseChecker:
     """Fixture to provide a hass_enforce_class_module checker."""
-    enforce_coordinator_module_checker = (
+    enforce_class_module_checker = hass_enforce_class_module.HassEnforceClassModule(
-        hass_enforce_class_module.HassEnforceClassModule(linter)
+        linter
     )
-    enforce_coordinator_module_checker.module = "homeassistant.components.pylint_test"
+    enforce_class_module_checker.module = "homeassistant.components.pylint_test"
-    return enforce_coordinator_module_checker
+    return enforce_class_module_checker

tests/pylint/test_enforce_class_module.py

@@ -41,21 +41,21 @@ from . import assert_adds_messages, assert_no_messages
         ),
     ],
 )
-def test_enforce_coordinator_module_good(
+def test_enforce_class_module_good(
-    linter: UnittestLinter, enforce_coordinator_module_checker: BaseChecker, code: str
+    linter: UnittestLinter, enforce_class_module_checker: BaseChecker, code: str
 ) -> None:
     """Good test cases."""
     root_node = astroid.parse(code, "homeassistant.components.pylint_test.coordinator")
     walker = ASTWalker(linter)
-    walker.add_checker(enforce_coordinator_module_checker)
+    walker.add_checker(enforce_class_module_checker)
 
     with assert_no_messages(linter):
         walker.walk(root_node)
 
 
-def test_enforce_coordinator_module_bad_simple(
+def test_enforce_class_module_bad_simple(
     linter: UnittestLinter,
-    enforce_coordinator_module_checker: BaseChecker,
+    enforce_class_module_checker: BaseChecker,
 ) -> None:
     """Bad test case with coordinator extending directly."""
     root_node = astroid.parse(
@@ -69,7 +69,7 @@ def test_enforce_coordinator_module_bad_simple(
         "homeassistant.components.pylint_test",
     )
     walker = ASTWalker(linter)
-    walker.add_checker(enforce_coordinator_module_checker)
+    walker.add_checker(enforce_class_module_checker)
 
     with assert_adds_messages(
         linter,
@@ -77,7 +77,7 @@ def test_enforce_coordinator_module_bad_simple(
             msg_id="hass-enforce-class-module",
             line=5,
             node=root_node.body[1],
-            args=None,
+            args=("DataUpdateCoordinator", "coordinator"),
             confidence=UNDEFINED,
             col_offset=0,
             end_line=5,
@@ -87,9 +87,9 @@ def test_enforce_coordinator_module_bad_simple(
         walker.walk(root_node)
 
 
-def test_enforce_coordinator_module_bad_nested(
+def test_enforce_class_module_bad_nested(
     linter: UnittestLinter,
-    enforce_coordinator_module_checker: BaseChecker,
+    enforce_class_module_checker: BaseChecker,
 ) -> None:
     """Bad test case with nested coordinators."""
     root_node = astroid.parse(
@@ -106,7 +106,7 @@ def test_enforce_coordinator_module_bad_nested(
         "homeassistant.components.pylint_test",
     )
     walker = ASTWalker(linter)
-    walker.add_checker(enforce_coordinator_module_checker)
+    walker.add_checker(enforce_class_module_checker)
 
     with assert_adds_messages(
         linter,
@@ -114,7 +114,7 @@ def test_enforce_coordinator_module_bad_nested(
             msg_id="hass-enforce-class-module",
             line=5,
             node=root_node.body[1],
-            args=None,
+            args=("DataUpdateCoordinator", "coordinator"),
             confidence=UNDEFINED,
             col_offset=0,
             end_line=5,
@@ -124,7 +124,7 @@ def test_enforce_coordinator_module_bad_nested(
             msg_id="hass-enforce-class-module",
             line=8,
             node=root_node.body[2],
-            args=None,
+            args=("DataUpdateCoordinator", "coordinator"),
             confidence=UNDEFINED,
             col_offset=0,
             end_line=8,