From 025c6f5e2e86e34948b9b1a2b0908220cf792d5a Mon Sep 17 00:00:00 2001
From: Sid <27780930+autinerd@users.noreply.github.com>
Date: Mon, 8 Apr 2024 20:44:18 +0200
Subject: [PATCH] Add `__slots__` to NodeClass classes (#115079)

Co-authored-by: J. Nick Koston <nick@koston.org>
---
 homeassistant/util/yaml/loader.py  | 31 +++++++++++++++---------------
 homeassistant/util/yaml/objects.py | 15 +++++++++++++++
 tests/test_config.py               |  1 -
 3 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/homeassistant/util/yaml/loader.py b/homeassistant/util/yaml/loader.py
index 79ee2797ad9..3a779b5e944 100644
--- a/homeassistant/util/yaml/loader.py
+++ b/homeassistant/util/yaml/loader.py
@@ -8,7 +8,7 @@ from io import StringIO, TextIOWrapper
 import logging
 import os
 from pathlib import Path
-from typing import Any, TextIO, TypeVar, overload
+from typing import Any, TextIO, overload
 
 import yaml
 
@@ -33,7 +33,6 @@ from .objects import Input, NodeDictClass, NodeListClass, NodeStrClass
 # mypy: allow-untyped-calls, no-warn-return-any
 
 JSON_TYPE = list | dict | str
-_DictT = TypeVar("_DictT", bound=dict)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -286,37 +285,37 @@ def _parse_yaml(
 
 @overload
 def _add_reference(
-    obj: list | NodeListClass,
-    loader: LoaderType,
-    node: yaml.nodes.Node,
+    obj: list | NodeListClass, loader: LoaderType, node: yaml.nodes.Node
 ) -> NodeListClass: ...
 
 
 @overload
 def _add_reference(
-    obj: str | NodeStrClass,
-    loader: LoaderType,
-    node: yaml.nodes.Node,
+    obj: str | NodeStrClass, loader: LoaderType, node: yaml.nodes.Node
 ) -> NodeStrClass: ...
 
 
 @overload
 def _add_reference(
-    obj: _DictT, loader: LoaderType, node: yaml.nodes.Node
-) -> _DictT: ...
+    obj: dict | NodeDictClass, loader: LoaderType, node: yaml.nodes.Node
+) -> NodeDictClass: ...
 
 
-def _add_reference(  # type: ignore[no-untyped-def]
-    obj, loader: LoaderType, node: yaml.nodes.Node
-):
+def _add_reference(
+    obj: dict | list | str | NodeDictClass | NodeListClass | NodeStrClass,
+    loader: LoaderType,
+    node: yaml.nodes.Node,
+) -> NodeDictClass | NodeListClass | NodeStrClass:
     """Add file reference information to an object."""
     if isinstance(obj, list):
         obj = NodeListClass(obj)
-    if isinstance(obj, str):
+    elif isinstance(obj, str):
         obj = NodeStrClass(obj)
+    elif isinstance(obj, dict):
+        obj = NodeDictClass(obj)
     try:  # suppress is much slower
-        setattr(obj, "__config_file__", loader.get_name)
-        setattr(obj, "__line__", node.start_mark.line + 1)
+        obj.__config_file__ = loader.get_name
+        obj.__line__ = node.start_mark.line + 1
     except AttributeError:
         pass
     return obj
diff --git a/homeassistant/util/yaml/objects.py b/homeassistant/util/yaml/objects.py
index 70c229c1a2f..d35ba11d25e 100644
--- a/homeassistant/util/yaml/objects.py
+++ b/homeassistant/util/yaml/objects.py
@@ -13,10 +13,20 @@ import yaml
 class NodeListClass(list):
     """Wrapper class to be able to add attributes on a list."""
 
+    __slots__ = ("__config_file__", "__line__")
+
+    __config_file__: str
+    __line__: int | str
+
 
 class NodeStrClass(str):
     """Wrapper class to be able to add attributes on a string."""
 
+    __slots__ = ("__config_file__", "__line__")
+
+    __config_file__: str
+    __line__: int | str
+
     def __voluptuous_compile__(self, schema: vol.Schema) -> Any:
         """Needed because vol.Schema.compile does not handle str subclasses."""
         return _compile_scalar(self)
@@ -25,6 +35,11 @@ class NodeStrClass(str):
 class NodeDictClass(dict):
     """Wrapper class to be able to add attributes on a dict."""
 
+    __slots__ = ("__config_file__", "__line__")
+
+    __config_file__: str
+    __line__: int | str
+
 
 @dataclass(slots=True, frozen=True)
 class Input:
diff --git a/tests/test_config.py b/tests/test_config.py
index 89a9b7b7082..defd6a1018b 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -381,7 +381,6 @@ async def mock_custom_validator_integrations_with_docs(
 class ConfigTestClass(NodeDictClass):
     """Test class for config with wrapper."""
 
-    __dict__ = {"__config_file__": "configuration.yaml", "__line__": 140}
     __line__ = 140
     __config_file__ = "configuration.yaml"