Rename YAML loader classes (#103609)

pull/103631/head
Erik Montnemery 2023-11-08 00:26:54 +01:00 committed by GitHub
parent 41a235bb52
commit e49f6b41ee
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 43 additions and 52 deletions

View File

@ -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 = "<unicode string>"
elif isinstance(stream, bytes):
self.name = "<byte string>"
else:
self.name = getattr(stream, "name", "<file>")
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 = "<unicode string>"
elif isinstance(stream, bytes):
self.name = "<byte string>"
else:
self.name = getattr(stream, "name", "<file>")
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)