From 05fc7cfbde573d55bf0a65741dbc486df444fd63 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Tue, 28 May 2024 18:15:53 +0200
Subject: [PATCH] Enforce namespace use for import conventions (#118215)

* Enforce namespace use for import conventions

* Include all registries

* Only apply to functions

* Use blacklist

* Rephrase comment

* Add async_entries_for_config_entry

* Typo

* Improve

* More core files

* Revert "More core files"

This reverts commit 9978b9370629af402a9a18f184b6f3a7ad45b08d.

* Revert diagnostics amends

* Include category/floor/label registries

* Performance

* Adjust text
---
 pylint/plugins/hass_imports.py | 52 +++++++++++++++++++++++++++++++
 tests/pylint/test_imports.py   | 57 ++++++++++++++++++++++++++++++++++
 2 files changed, 109 insertions(+)

diff --git a/pylint/plugins/hass_imports.py b/pylint/plugins/hass_imports.py
index d8f85df011f..b4d30be483d 100644
--- a/pylint/plugins/hass_imports.py
+++ b/pylint/plugins/hass_imports.py
@@ -395,6 +395,38 @@ _OBSOLETE_IMPORT: dict[str, list[ObsoleteImportMatch]] = {
 }
 
 
+# Blacklist of imports that should be using the namespace
+@dataclass
+class NamespaceAlias:
+    """Class for namespace imports."""
+
+    alias: str
+    names: set[str]  # function names
+
+
+_FORCE_NAMESPACE_IMPORT: dict[str, NamespaceAlias] = {
+    "homeassistant.helpers.area_registry": NamespaceAlias("ar", {"async_get"}),
+    "homeassistant.helpers.category_registry": NamespaceAlias("cr", {"async_get"}),
+    "homeassistant.helpers.device_registry": NamespaceAlias(
+        "dr",
+        {
+            "async_get",
+            "async_entries_for_config_entry",
+        },
+    ),
+    "homeassistant.helpers.entity_registry": NamespaceAlias(
+        "er",
+        {
+            "async_get",
+            "async_entries_for_config_entry",
+        },
+    ),
+    "homeassistant.helpers.floor_registry": NamespaceAlias("fr", {"async_get"}),
+    "homeassistant.helpers.issue_registry": NamespaceAlias("ir", {"async_get"}),
+    "homeassistant.helpers.label_registry": NamespaceAlias("lr", {"async_get"}),
+}
+
+
 class HassImportsFormatChecker(BaseChecker):
     """Checker for imports."""
 
@@ -422,6 +454,12 @@ class HassImportsFormatChecker(BaseChecker):
             "Used when an import from another component should be "
             "from the component root",
         ),
+        "W7425": (
+            "`%s` should not be imported directly. Please import `%s` as `%s` "
+            "and use `%s.%s`",
+            "hass-helper-namespace-import",
+            "Used when a helper should be used via the namespace",
+        ),
     }
     options = ()
 
@@ -524,6 +562,20 @@ class HassImportsFormatChecker(BaseChecker):
                             node=node,
                             args=(import_match.string, obsolete_import.reason),
                         )
+        if namespace_alias := _FORCE_NAMESPACE_IMPORT.get(node.modname):
+            for name in node.names:
+                if name[0] in namespace_alias.names:
+                    self.add_message(
+                        "hass-helper-namespace-import",
+                        node=node,
+                        args=(
+                            name[0],
+                            node.modname,
+                            namespace_alias.alias,
+                            namespace_alias.alias,
+                            name[0],
+                        ),
+                    )
 
 
 def register(linter: PyLinter) -> None:
diff --git a/tests/pylint/test_imports.py b/tests/pylint/test_imports.py
index 5f1d4d86840..e53b8206848 100644
--- a/tests/pylint/test_imports.py
+++ b/tests/pylint/test_imports.py
@@ -252,3 +252,60 @@ def test_bad_root_import(
             imports_checker.visit_import(node)
         if import_node.startswith("from"):
             imports_checker.visit_importfrom(node)
+
+
+@pytest.mark.parametrize(
+    ("import_node", "module_name", "expected_args"),
+    [
+        (
+            "from homeassistant.helpers.issue_registry import async_get",
+            "tests.components.pylint_test.climate",
+            (
+                "async_get",
+                "homeassistant.helpers.issue_registry",
+                "ir",
+                "ir",
+                "async_get",
+            ),
+        ),
+        (
+            "from homeassistant.helpers.issue_registry import async_get as async_get_issue_registry",
+            "tests.components.pylint_test.climate",
+            (
+                "async_get",
+                "homeassistant.helpers.issue_registry",
+                "ir",
+                "ir",
+                "async_get",
+            ),
+        ),
+    ],
+)
+def test_bad_namespace_import(
+    linter: UnittestLinter,
+    imports_checker: BaseChecker,
+    import_node: str,
+    module_name: str,
+    expected_args: tuple[str, ...],
+) -> None:
+    """Ensure bad namespace imports are rejected."""
+
+    node = astroid.extract_node(
+        f"{import_node} #@",
+        module_name,
+    )
+    imports_checker.visit_module(node.parent)
+
+    with assert_adds_messages(
+        linter,
+        pylint.testutils.MessageTest(
+            msg_id="hass-helper-namespace-import",
+            node=node,
+            args=expected_args,
+            line=1,
+            col_offset=0,
+            end_line=1,
+            end_col_offset=len(import_node),
+        ),
+    ):
+        imports_checker.visit_importfrom(node)