Add __slots__ to NodeClass classes (#115079)

Co-authored-by: J. Nick Koston <nick@koston.org>

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

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:

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"