Add __slots__ to NodeClass classes (#115079)

Co-authored-by: J. Nick Koston <nick@koston.org>
This commit is contained in:
Sid 2024-04-08 20:44:18 +02:00 committed by GitHub
parent 266e4f0c8b
commit 025c6f5e2e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 30 additions and 17 deletions

View File

@ -8,7 +8,7 @@ from io import StringIO, TextIOWrapper
import logging import logging
import os import os
from pathlib import Path from pathlib import Path
from typing import Any, TextIO, TypeVar, overload from typing import Any, TextIO, overload
import yaml import yaml
@ -33,7 +33,6 @@ from .objects import Input, NodeDictClass, NodeListClass, NodeStrClass
# mypy: allow-untyped-calls, no-warn-return-any # mypy: allow-untyped-calls, no-warn-return-any
JSON_TYPE = list | dict | str JSON_TYPE = list | dict | str
_DictT = TypeVar("_DictT", bound=dict)
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
@ -286,37 +285,37 @@ def _parse_yaml(
@overload @overload
def _add_reference( def _add_reference(
obj: list | NodeListClass, obj: list | NodeListClass, loader: LoaderType, node: yaml.nodes.Node
loader: LoaderType,
node: yaml.nodes.Node,
) -> NodeListClass: ... ) -> NodeListClass: ...
@overload @overload
def _add_reference( def _add_reference(
obj: str | NodeStrClass, obj: str | NodeStrClass, loader: LoaderType, node: yaml.nodes.Node
loader: LoaderType,
node: yaml.nodes.Node,
) -> NodeStrClass: ... ) -> NodeStrClass: ...
@overload @overload
def _add_reference( def _add_reference(
obj: _DictT, loader: LoaderType, node: yaml.nodes.Node obj: dict | NodeDictClass, loader: LoaderType, node: yaml.nodes.Node
) -> _DictT: ... ) -> NodeDictClass: ...
def _add_reference( # type: ignore[no-untyped-def] def _add_reference(
obj, loader: LoaderType, node: yaml.nodes.Node obj: dict | list | str | NodeDictClass | NodeListClass | NodeStrClass,
): loader: LoaderType,
node: yaml.nodes.Node,
) -> NodeDictClass | NodeListClass | NodeStrClass:
"""Add file reference information to an object.""" """Add file reference information to an object."""
if isinstance(obj, list): if isinstance(obj, list):
obj = NodeListClass(obj) obj = NodeListClass(obj)
if isinstance(obj, str): elif isinstance(obj, str):
obj = NodeStrClass(obj) obj = NodeStrClass(obj)
elif isinstance(obj, dict):
obj = NodeDictClass(obj)
try: # suppress is much slower try: # suppress is much slower
setattr(obj, "__config_file__", loader.get_name) obj.__config_file__ = loader.get_name
setattr(obj, "__line__", node.start_mark.line + 1) obj.__line__ = node.start_mark.line + 1
except AttributeError: except AttributeError:
pass pass
return obj return obj

View File

@ -13,10 +13,20 @@ import yaml
class NodeListClass(list): class NodeListClass(list):
"""Wrapper class to be able to add attributes on a 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): class NodeStrClass(str):
"""Wrapper class to be able to add attributes on a string.""" """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: def __voluptuous_compile__(self, schema: vol.Schema) -> Any:
"""Needed because vol.Schema.compile does not handle str subclasses.""" """Needed because vol.Schema.compile does not handle str subclasses."""
return _compile_scalar(self) return _compile_scalar(self)
@ -25,6 +35,11 @@ class NodeStrClass(str):
class NodeDictClass(dict): class NodeDictClass(dict):
"""Wrapper class to be able to add attributes on a 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) @dataclass(slots=True, frozen=True)
class Input: class Input:

View File

@ -381,7 +381,6 @@ async def mock_custom_validator_integrations_with_docs(
class ConfigTestClass(NodeDictClass): class ConfigTestClass(NodeDictClass):
"""Test class for config with wrapper.""" """Test class for config with wrapper."""
__dict__ = {"__config_file__": "configuration.yaml", "__line__": 140}
__line__ = 140 __line__ = 140
__config_file__ = "configuration.yaml" __config_file__ = "configuration.yaml"