Split up yaml loaders into multiple files

homeassistant/scripts/check_config.py

@@ -17,20 +17,23 @@
     CONF_PACKAGES, merge_packages_config, _format_config_error,
     find_config_file, load_yaml_config_file,
     extract_domain_configs, config_per_platform)
-from homeassistant.util import yaml
+
+import homeassistant.util.yaml.loader as yaml_loader
 from homeassistant.exceptions import HomeAssistantError
 
 REQUIREMENTS = ('colorlog==4.0.2',)
 
 _LOGGER = logging.getLogger(__name__)
 # pylint: disable=protected-access
 MOCKS = {
-    'load': ("homeassistant.util.yaml.load_yaml", yaml.load_yaml),
-    'load*': ("homeassistant.config.load_yaml", yaml.load_yaml),
-    'secrets': ("homeassistant.util.yaml.secret_yaml", yaml.secret_yaml),
+    'load': ("homeassistant.util.yaml.loader.load_yaml",
+             yaml_loader.load_yaml),
+    'load*': ("homeassistant.config.load_yaml", yaml_loader.load_yaml),
+    'secrets': ("homeassistant.util.yaml.loader.secret_yaml",
+                yaml_loader.secret_yaml),
 }
 SILENCE = (
-    'homeassistant.scripts.check_config.yaml.clear_secret_cache',
+    'homeassistant.scripts.check_config.yaml_loader.clear_secret_cache',
 )
 
 PATCHES = {}
@@ -195,15 +198,16 @@ def mock_secrets(ldr, node):
 
     if secrets:
         # Ensure !secrets point to the patched function
-        yaml.yaml.SafeLoader.add_constructor('!secret', yaml.secret_yaml)
+        yaml_loader.yaml.SafeLoader.add_constructor('!secret',
+                                                    yaml_loader.secret_yaml)
 
     try:
         hass = core.HomeAssistant()
         hass.config.config_dir = config_dir
 
         res['components'] = hass.loop.run_until_complete(
             check_ha_config_file(hass))
-        res['secret_cache'] = OrderedDict(yaml.__SECRET_CACHE)
+        res['secret_cache'] = OrderedDict(yaml_loader.__SECRET_CACHE)
 
         for err in res['components'].errors:
             domain = err.domain or ERROR_STR
@@ -221,7 +225,8 @@ def mock_secrets(ldr, node):
             pat.stop()
         if secrets:
             # Ensure !secrets point to the original function
-            yaml.yaml.SafeLoader.add_constructor('!secret', yaml.secret_yaml)
+            yaml_loader.yaml.SafeLoader.add_constructor(
+                '!secret', yaml_loader.secret_yaml)
         bootstrap.clear_secret_cache()
 
     return res
@@ -239,7 +244,7 @@ def line_info(obj, **kwargs):
 def dump_dict(layer, indent_count=3, listi=False, **kwargs):
     """Display a dict.
 
-    A friendly version of print yaml.yaml.dump(config).
+    A friendly version of print yaml_loader.yaml.dump(config).
     """
     def sort_dict_key(val):
         """Return the dict key for sorting."""
@@ -311,7 +316,7 @@ def _comp_error(ex, domain, config):
         return result.add_error(
             "Error loading {}: {}".format(config_path, err))
     finally:
-        yaml.clear_secret_cache()
+        yaml_loader.clear_secret_cache()
 
     # Extract and validate core [homeassistant] config
     try:

homeassistant/util/yaml/__init__.py

@@ -0,0 +1,15 @@
+"""YAML utility functions."""
+from .const import (
+    SECRET_YAML, _SECRET_NAMESPACE
+)
+from .dumper import dump, save_yaml
+from .loader import (
+    clear_secret_cache, load_yaml, secret_yaml
+)
+
+
+__all__ = [
+    'SECRET_YAML', '_SECRET_NAMESPACE',
+    'dump', 'save_yaml',
+    'clear_secret_cache', 'load_yaml', 'secret_yaml',
+]

homeassistant/util/yaml/const.py

@@ -0,0 +1,4 @@
+"""Constants."""
+SECRET_YAML = 'secrets.yaml'
+
+_SECRET_NAMESPACE = 'homeassistant'

homeassistant/util/yaml/dumper.py

@@ -0,0 +1,60 @@
+"""Custom dumper and representers."""
+from collections import OrderedDict
+import yaml
+
+from .objects import NodeListClass
+
+
+def dump(_dict: dict) -> str:
+    """Dump YAML to a string and remove null."""
+    return yaml.safe_dump(
+        _dict, default_flow_style=False, allow_unicode=True) \
+        .replace(': null\n', ':\n')
+
+
+def save_yaml(path: str, data: dict) -> None:
+    """Save YAML to a file."""
+    # Dump before writing to not truncate the file if dumping fails
+    str_data = dump(data)
+    with open(path, 'w', encoding='utf-8') as outfile:
+        outfile.write(str_data)
+
+
+# From: https://gist.github.com/miracle2k/3184458
+# pylint: disable=redefined-outer-name
+def represent_odict(dump, tag, mapping,  # type: ignore
+                    flow_style=None) -> yaml.MappingNode:
+    """Like BaseRepresenter.represent_mapping but does not issue the sort()."""
+    value = []  # type: list
+    node = yaml.MappingNode(tag, value, flow_style=flow_style)
+    if dump.alias_key is not None:
+        dump.represented_objects[dump.alias_key] = node
+    best_style = True
+    if hasattr(mapping, 'items'):
+        mapping = mapping.items()
+    for item_key, item_value in mapping:
+        node_key = dump.represent_data(item_key)
+        node_value = dump.represent_data(item_value)
+        if not (isinstance(node_key, yaml.ScalarNode) and not node_key.style):
+            best_style = False
+        if not (isinstance(node_value, yaml.ScalarNode) and
+                not node_value.style):
+            best_style = False
+        value.append((node_key, node_value))
+    if flow_style is None:
+        if dump.default_flow_style is not None:
+            node.flow_style = dump.default_flow_style
+        else:
+            node.flow_style = best_style
+    return node
+
+
+yaml.SafeDumper.add_representer(
+    OrderedDict,
+    lambda dumper, value:
+    represent_odict(dumper, 'tag:yaml.org,2002:map', value))
+
+yaml.SafeDumper.add_representer(
+    NodeListClass,
+    lambda dumper, value:
+    dumper.represent_sequence('tag:yaml.org,2002:seq', value))

homeassistant/util/yaml.py → homeassistant/util/yaml/loader.py

@@ -1,4 +1,4 @@
-"""YAML utility functions."""
+"""Custom loader."""
 import logging
 import os
 import sys
@@ -7,6 +7,7 @@
 from typing import Union, List, Dict, Iterator, overload, TypeVar
 
 import yaml
+
 try:
     import keyring
 except ImportError:
@@ -19,25 +20,23 @@
 
 from homeassistant.exceptions import HomeAssistantError
 
+from .const import _SECRET_NAMESPACE, SECRET_YAML
+from .objects import NodeListClass, NodeStrClass
+
+
 _LOGGER = logging.getLogger(__name__)
-_SECRET_NAMESPACE = 'homeassistant'
-SECRET_YAML = 'secrets.yaml'
 __SECRET_CACHE = {}  # type: Dict[str, JSON_TYPE]
 
 JSON_TYPE = Union[List, Dict, str]  # pylint: disable=invalid-name
 DICT_T = TypeVar('DICT_T', bound=Dict)  # pylint: disable=invalid-name
 
 
-class NodeListClass(list):
-    """Wrapper class to be able to add attributes on a list."""
-
-    pass
-
-
-class NodeStrClass(str):
-    """Wrapper class to be able to add attributes on a string."""
+def clear_secret_cache() -> None:
+    """Clear the secret cache.
 
-    pass
+    Async friendly.
+    """
+    __SECRET_CACHE.clear()
 
 
 # pylint: disable=too-many-ancestors
@@ -54,6 +53,21 @@ def compose_node(self, parent: yaml.nodes.Node,
         return node
 
 
+def load_yaml(fname: str) -> JSON_TYPE:
+    """Load a YAML file."""
+    try:
+        with open(fname, encoding='utf-8') as conf_file:
+            # If configuration file is empty YAML returns None
+            # We convert that to an empty dict
+            return yaml.load(conf_file, Loader=SafeLineLoader) or OrderedDict()
+    except yaml.YAMLError as exc:
+        _LOGGER.error(str(exc))
+        raise HomeAssistantError(exc)
+    except UnicodeDecodeError as exc:
+        _LOGGER.error("Unable to read file %s: %s", fname, exc)
+        raise HomeAssistantError(exc)
+
+
 # pylint: disable=pointless-statement
 @overload
 def _add_reference(obj: Union[list, NodeListClass],
@@ -86,44 +100,6 @@ def _add_reference(obj, loader: SafeLineLoader,  # type: ignore # noqa: F811
     return obj
 
 
-def load_yaml(fname: str) -> JSON_TYPE:
-    """Load a YAML file."""
-    try:
-        with open(fname, encoding='utf-8') as conf_file:
-            # If configuration file is empty YAML returns None
-            # We convert that to an empty dict
-            return yaml.load(conf_file, Loader=SafeLineLoader) or OrderedDict()
-    except yaml.YAMLError as exc:
-        _LOGGER.error(str(exc))
-        raise HomeAssistantError(exc)
-    except UnicodeDecodeError as exc:
-        _LOGGER.error("Unable to read file %s: %s", fname, exc)
-        raise HomeAssistantError(exc)
-
-
-def dump(_dict: dict) -> str:
-    """Dump YAML to a string and remove null."""
-    return yaml.safe_dump(
-        _dict, default_flow_style=False, allow_unicode=True) \
-        .replace(': null\n', ':\n')
-
-
-def save_yaml(path: str, data: dict) -> None:
-    """Save YAML to a file."""
-    # Dump before writing to not truncate the file if dumping fails
-    str_data = dump(data)
-    with open(path, 'w', encoding='utf-8') as outfile:
-        outfile.write(str_data)
-
-
-def clear_secret_cache() -> None:
-    """Clear the secret cache.
-
-    Async friendly.
-    """
-    __SECRET_CACHE.clear()
-
-
 def _include_yaml(loader: SafeLineLoader,
                   node: yaml.nodes.Node) -> JSON_TYPE:
     """Load another YAML file and embeds it using the !include tag.
@@ -331,43 +307,3 @@ def secret_yaml(loader: SafeLineLoader,
 yaml.SafeLoader.add_constructor('!include_dir_named', _include_dir_named_yaml)
 yaml.SafeLoader.add_constructor('!include_dir_merge_named',
                                 _include_dir_merge_named_yaml)
-
-
-# From: https://gist.github.com/miracle2k/3184458
-# pylint: disable=redefined-outer-name
-def represent_odict(dump, tag, mapping,  # type: ignore
-                    flow_style=None) -> yaml.MappingNode:
-    """Like BaseRepresenter.represent_mapping but does not issue the sort()."""
-    value = []  # type: list
-    node = yaml.MappingNode(tag, value, flow_style=flow_style)
-    if dump.alias_key is not None:
-        dump.represented_objects[dump.alias_key] = node
-    best_style = True
-    if hasattr(mapping, 'items'):
-        mapping = mapping.items()
-    for item_key, item_value in mapping:
-        node_key = dump.represent_data(item_key)
-        node_value = dump.represent_data(item_value)
-        if not (isinstance(node_key, yaml.ScalarNode) and not node_key.style):
-            best_style = False
-        if not (isinstance(node_value, yaml.ScalarNode) and
-                not node_value.style):
-            best_style = False
-        value.append((node_key, node_value))
-    if flow_style is None:
-        if dump.default_flow_style is not None:
-            node.flow_style = dump.default_flow_style
-        else:
-            node.flow_style = best_style
-    return node
-
-
-yaml.SafeDumper.add_representer(
-    OrderedDict,
-    lambda dumper, value:
-    represent_odict(dumper, 'tag:yaml.org,2002:map', value))
-
-yaml.SafeDumper.add_representer(
-    NodeListClass,
-    lambda dumper, value:
-    dumper.represent_sequence('tag:yaml.org,2002:seq', value))

homeassistant/util/yaml/objects.py

@@ -0,0 +1,13 @@
+"""Custom yaml object types."""
+
+
+class NodeListClass(list):
+    """Wrapper class to be able to add attributes on a list."""
+
+    pass
+
+
+class NodeStrClass(str):
+    """Wrapper class to be able to add attributes on a string."""
+
+    pass

mypy.ini

@@ -17,7 +17,11 @@ disallow_untyped_defs = true
 [mypy-homeassistant.config_entries]
 disallow_untyped_defs = false
 
-[mypy-homeassistant.util.yaml]
+[mypy-homeassistant.util.yaml.dumper]
+warn_return_any = false
+disallow_untyped_calls = false
+
+[mypy-homeassistant.util.yaml.loader]
 warn_return_any = false
 disallow_untyped_calls = false
 

tests/common.py

@@ -15,7 +15,8 @@
 from unittest.mock import MagicMock, Mock, patch
 
 import homeassistant.util.dt as date_util
-import homeassistant.util.yaml as yaml
+import homeassistant.util.yaml.loader as yaml_loader
+import homeassistant.util.yaml.dumper as yaml_dumper
 
 from homeassistant import auth, config_entries, core as ha, loader
 from homeassistant.auth import (
@@ -680,7 +681,8 @@ def mock_open_f(fname, **_):
         # Not found
         raise FileNotFoundError("File not found: {}".format(fname))
 
-    return patch.object(yaml, 'open', mock_open_f, create=True)
+    return patch.object(yaml_loader, 'open', mock_open_f, create=True)
+    return patch.object(yaml_dumper, 'open', mock_open_f, create=True)
 
 
 def mock_coro(return_value=None, exception=None):

tests/components/scene/test_init.py

@@ -4,7 +4,7 @@
 
 from homeassistant.setup import setup_component
 from homeassistant.components import light, scene
-from homeassistant.util import yaml
+from homeassistant.util.yaml import loader as yaml_loader
 
 from tests.common import get_test_home_assistant
 from tests.components.light import common as common_light
@@ -90,7 +90,7 @@ def test_config_yaml_bool(self):
                 self.light_1.entity_id, self.light_2.entity_id)
 
         with io.StringIO(config) as file:
-            doc = yaml.yaml.safe_load(file)
+            doc = yaml_loader.yaml.load(file)
 
         assert setup_component(self.hass, scene.DOMAIN, doc)
         common.activate(self.hass, 'scene.test')

tests/helpers/test_entity_registry.py

@@ -11,7 +11,7 @@
 from tests.common import mock_registry, flush_store
 
 
-YAML__OPEN_PATH = 'homeassistant.util.yaml.open'
+YAML__OPEN_PATH = 'homeassistant.util.yaml.loader.open'
 
 
 @pytest.fixture