diff --git a/homeassistant/components/dsmr_reader/definitions.py b/homeassistant/components/dsmr_reader/definitions.py
index 9003c4d4334d40d1b5b477a4bd932d829c9f9fab..62d095aa993ff1949fe7d41022fedec53696246b 100644
--- a/homeassistant/components/dsmr_reader/definitions.py
+++ b/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."""
diff --git a/homeassistant/components/growatt_server/sensor_types/sensor_entity_description.py b/homeassistant/components/growatt_server/sensor_types/sensor_entity_description.py
index e1ee4c303267a9887bc3b1a2307c475fc46eafe2..10d00671ba52a4aab9288e58c0e189f45ffe77dd 100644
--- a/homeassistant/components/growatt_server/sensor_types/sensor_entity_description.py
+++ b/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."""
diff --git a/homeassistant/components/repetier/__init__.py b/homeassistant/components/repetier/__init__.py
index 2642e78e7ec0b2be1b66a7b33e4b59440e176441..27ddc62a8471042a9e8ff46b94f0bd8e5df4e4d1 100644
--- a/homeassistant/components/repetier/__init__.py
+++ b/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
):
diff --git a/homeassistant/components/sunweg/sensor_types/sensor_entity_description.py b/homeassistant/components/sunweg/sensor_types/sensor_entity_description.py
index 8c792ab617fd8cb5b809dc10af76b7f99691c140..1d06f04ab3d4dc41c11805c275541ac855dfe387 100644
--- a/homeassistant/components/sunweg/sensor_types/sensor_entity_description.py
+++ b/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."""
diff --git a/pylint/plugins/hass_enforce_class_module.py b/pylint/plugins/hass_enforce_class_module.py
index d9f844f907f0e4eae0c5cbf2d2b6a9506ce80540..dcd42f9a1c16c78c2c2aaf5a5c9fec6315f31ef7 100644
--- a/pylint/plugins/hass_enforce_class_module.py
+++ b/pylint/plugins/hass_enforce_class_module.py
@@ -1,38 +1,91 @@
-"""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
- if not root_name.startswith("homeassistant.components"):
+ # we only want to check components
+ if not root_name.startswith("homeassistant.components."):
return
- is_coordinator_module = root_name.endswith(".coordinator")
- for ancestor in node.ancestors():
- if ancestor.name == "DataUpdateCoordinator" and not is_coordinator_module:
- self.add_message("hass-enforce-class-module", node=node)
- return
+ ancestors: list[ClassDef] | None = None
+
+ for match in _MODULES:
+ 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
def register(linter: PyLinter) -> None:
diff --git a/tests/pylint/conftest.py b/tests/pylint/conftest.py
index 38b4188230f1c5e043ae845b2a9bac86d84e437c..5e8ed28da6b477258c02808e5dd023a376069c5a 100644
--- a/tests/pylint/conftest.py
+++ b/tests/pylint/conftest.py
@@ -113,13 +113,11 @@ def hass_enforce_class_module_fixture() -> ModuleType:
)
-@pytest.fixture(name="enforce_coordinator_module_checker")
-def enforce_coordinator_module_fixture(
- hass_enforce_class_module, linter
-) -> BaseChecker:
+@pytest.fixture(name="enforce_class_module_checker")
+def enforce_class_module_fixture(hass_enforce_class_module, linter) -> BaseChecker:
"""Fixture to provide a hass_enforce_class_module checker."""
- enforce_coordinator_module_checker = (
- hass_enforce_class_module.HassEnforceClassModule(linter)
+ enforce_class_module_checker = hass_enforce_class_module.HassEnforceClassModule(
+ linter
)
- enforce_coordinator_module_checker.module = "homeassistant.components.pylint_test"
- return enforce_coordinator_module_checker
+ enforce_class_module_checker.module = "homeassistant.components.pylint_test"
+ return enforce_class_module_checker
diff --git a/tests/pylint/test_enforce_class_module.py b/tests/pylint/test_enforce_class_module.py
index 5fd6e0e88ccd29737026b1481b3515863a76d2fa..b0f071fde52e209dc8abd955cd4bf2392c6776d9 100644
--- a/tests/pylint/test_enforce_class_module.py
+++ b/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(
- linter: UnittestLinter, enforce_coordinator_module_checker: BaseChecker, code: str
+def test_enforce_class_module_good(
+ 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,