diff --git a/homeassistant/util/yaml/loader.py b/homeassistant/util/yaml/loader.py index 73e7861902f..8a8822ab17f 100644 --- a/homeassistant/util/yaml/loader.py +++ b/homeassistant/util/yaml/loader.py @@ -101,48 +101,11 @@ class Secrets: return secrets -class SafeLoader(FastestAvailableSafeLoader): - """The fastest available safe loader.""" +class _LoaderMixin: + """Mixin class with extensions for YAML loader.""" - def __init__(self, stream: Any, secrets: Secrets | None = None) -> None: - """Initialize a safe line loader.""" - self.stream = stream - if isinstance(stream, str): - self.name = "" - elif isinstance(stream, bytes): - self.name = "" - else: - self.name = getattr(stream, "name", "") - super().__init__(stream) - self.secrets = secrets - - def get_name(self) -> str: - """Get the name of the loader.""" - return self.name - - def get_stream_name(self) -> str: - """Get the name of the stream.""" - return self.stream.name or "" - - -class SafeLineLoader(yaml.SafeLoader): - """Loader class that keeps track of line numbers.""" - - def __init__(self, stream: Any, secrets: Secrets | None = None) -> None: - """Initialize a safe line loader.""" - super().__init__(stream) - self.secrets = secrets - - def compose_node( # type: ignore[override] - self, parent: yaml.nodes.Node, index: int - ) -> yaml.nodes.Node: - """Annotate a node with the first line it was seen.""" - last_line: int = self.line - node: yaml.nodes.Node = super().compose_node( # type: ignore[assignment] - parent, index - ) - node.__line__ = last_line + 1 # type: ignore[attr-defined] - return node + name: str + stream: Any def get_name(self) -> str: """Get the name of the loader.""" @@ -153,7 +116,35 @@ class SafeLineLoader(yaml.SafeLoader): return getattr(self.stream, "name", "") -LoaderType = SafeLineLoader | SafeLoader +class FastSafeLoader(FastestAvailableSafeLoader, _LoaderMixin): + """The fastest available safe loader, either C or Python.""" + + def __init__(self, stream: Any, secrets: Secrets | None = None) -> None: + """Initialize a safe line loader.""" + self.stream = stream + + # Set name in same way as the Python loader does in yaml.reader.__init__ + if isinstance(stream, str): + self.name = "" + elif isinstance(stream, bytes): + self.name = "" + else: + self.name = getattr(stream, "name", "") + + super().__init__(stream) + self.secrets = secrets + + +class PythonSafeLoader(yaml.SafeLoader, _LoaderMixin): + """Python safe loader.""" + + def __init__(self, stream: Any, secrets: Secrets | None = None) -> None: + """Initialize a safe line loader.""" + super().__init__(stream) + self.secrets = secrets + + +LoaderType = FastSafeLoader | PythonSafeLoader def load_yaml(fname: str, secrets: Secrets | None = None) -> JSON_TYPE: @@ -171,31 +162,31 @@ def parse_yaml( ) -> JSON_TYPE: """Parse YAML with the fastest available loader.""" if not HAS_C_LOADER: - return _parse_yaml_pure_python(content, secrets) + return _parse_yaml_python(content, secrets) try: - return _parse_yaml(SafeLoader, content, secrets) + return _parse_yaml(FastSafeLoader, content, secrets) except yaml.YAMLError: - # Loading failed, so we now load with the slow line loader - # since the C one will not give us line numbers + # Loading failed, so we now load with the Python loader which has more + # readable exceptions if isinstance(content, (StringIO, TextIO, TextIOWrapper)): # Rewind the stream so we can try again content.seek(0, 0) - return _parse_yaml_pure_python(content, secrets) + return _parse_yaml_python(content, secrets) -def _parse_yaml_pure_python( +def _parse_yaml_python( content: str | TextIO | StringIO, secrets: Secrets | None = None ) -> JSON_TYPE: - """Parse YAML with the pure python loader (this is very slow).""" + """Parse YAML with the python loader (this is very slow).""" try: - return _parse_yaml(SafeLineLoader, content, secrets) + return _parse_yaml(PythonSafeLoader, content, secrets) except yaml.YAMLError as exc: _LOGGER.error(str(exc)) raise HomeAssistantError(exc) from exc def _parse_yaml( - loader: type[SafeLoader] | type[SafeLineLoader], + loader: type[FastSafeLoader] | type[PythonSafeLoader], content: str | TextIO, secrets: Secrets | None = None, ) -> JSON_TYPE: @@ -404,7 +395,7 @@ def secret_yaml(loader: LoaderType, node: yaml.nodes.Node) -> JSON_TYPE: def add_constructor(tag: Any, constructor: Any) -> None: """Add to constructor to all loaders.""" - for yaml_loader in (SafeLoader, SafeLineLoader): + for yaml_loader in (FastSafeLoader, PythonSafeLoader): yaml_loader.add_constructor(tag, constructor)