Fix util.ruamel_yaml type errors (#25338)

This commit is contained in:
Ville Skyttä 2019-07-21 00:35:22 +03:00 committed by Paulus Schoutsen
parent 7ea27c0f2a
commit 56e4a2aea6

View File

@ -3,7 +3,7 @@ import logging
import os import os
from os import O_CREAT, O_TRUNC, O_WRONLY, stat_result from os import O_CREAT, O_TRUNC, O_WRONLY, stat_result
from collections import OrderedDict from collections import OrderedDict
from typing import Union, List, Dict from typing import Union, List, Dict, Optional
import ruamel.yaml import ruamel.yaml
from ruamel.yaml import YAML from ruamel.yaml import YAML
@ -22,6 +22,8 @@ JSON_TYPE = Union[List, Dict, str] # pylint: disable=invalid-name
class ExtSafeConstructor(SafeConstructor): class ExtSafeConstructor(SafeConstructor):
"""Extended SafeConstructor.""" """Extended SafeConstructor."""
name = None # type: Optional[str]
class UnsupportedYamlError(HomeAssistantError): class UnsupportedYamlError(HomeAssistantError):
"""Unsupported YAML.""" """Unsupported YAML."""
@ -31,22 +33,25 @@ class WriteError(HomeAssistantError):
"""Error writing the data.""" """Error writing the data."""
def _include_yaml(constructor: SafeConstructor, node: ruamel.yaml.nodes.Node) \ def _include_yaml(constructor: ExtSafeConstructor,
-> JSON_TYPE: node: ruamel.yaml.nodes.Node) -> JSON_TYPE:
"""Load another YAML file and embeds it using the !include tag. """Load another YAML file and embeds it using the !include tag.
Example: Example:
device_tracker: !include device_tracker.yaml device_tracker: !include device_tracker.yaml
""" """
if constructor.name is None:
raise HomeAssistantError(
"YAML include error: filename not set for %s" % node.value)
fname = os.path.join(os.path.dirname(constructor.name), node.value) fname = os.path.join(os.path.dirname(constructor.name), node.value)
return load_yaml(fname, False) return load_yaml(fname, False)
def _yaml_unsupported(constructor: SafeConstructor, node: def _yaml_unsupported(constructor: ExtSafeConstructor, node:
ruamel.yaml.nodes.Node) -> None: ruamel.yaml.nodes.Node) -> None:
raise UnsupportedYamlError( raise UnsupportedYamlError(
'Unsupported YAML, you can not use {} in {}' 'Unsupported YAML, you can not use {} in {}'
.format(node.tag, os.path.basename(constructor.name))) .format(node.tag, os.path.basename(constructor.name or '(None)')))
def object_to_yaml(data: JSON_TYPE) -> str: def object_to_yaml(data: JSON_TYPE) -> str:
@ -80,7 +85,7 @@ def load_yaml(fname: str, round_trip: bool = False) -> JSON_TYPE:
yaml = YAML(typ='rt') yaml = YAML(typ='rt')
yaml.preserve_quotes = True yaml.preserve_quotes = True
else: else:
if not hasattr(ExtSafeConstructor, 'name'): if ExtSafeConstructor.name is None:
ExtSafeConstructor.name = fname ExtSafeConstructor.name = fname
yaml = YAML(typ='safe') yaml = YAML(typ='safe')
yaml.Constructor = ExtSafeConstructor yaml.Constructor = ExtSafeConstructor