From da2521a29922c7da744685255f7f8d16fcd0cc1e Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <paulus@paulusschoutsen.nl>
Date: Mon, 1 May 2017 20:09:49 -0700
Subject: [PATCH] Fix YAML dump (#7388)

* Fix YAML dump

* Add test
---
 homeassistant/util/yaml.py | 29 ++++++++++++++++++-----------
 tests/util/test_yaml.py    |  8 ++++++++
 2 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/homeassistant/util/yaml.py b/homeassistant/util/yaml.py
index 19e614da1f8..7827f484fdf 100644
--- a/homeassistant/util/yaml.py
+++ b/homeassistant/util/yaml.py
@@ -20,18 +20,20 @@ _SECRET_YAML = 'secrets.yaml'
 __SECRET_CACHE = {}  # type: Dict
 
 
+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
+
+
 def _add_reference(obj, loader, node):
     """Add file reference information to an object."""
-    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
-
     if isinstance(obj, list):
         obj = NodeListClass(obj)
     if isinstance(obj, str):
@@ -305,4 +307,9 @@ def represent_odict(dump, tag, mapping, flow_style=None):
 yaml.SafeDumper.add_representer(
     OrderedDict,
     lambda dumper, value:
-    represent_odict(dumper, u'tag:yaml.org,2002:map', 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))
diff --git a/tests/util/test_yaml.py b/tests/util/test_yaml.py
index 79fd994ce86..0ccb6f5d6d0 100644
--- a/tests/util/test_yaml.py
+++ b/tests/util/test_yaml.py
@@ -385,3 +385,11 @@ class TestSecrets(unittest.TestCase):
         load_yaml(self._yaml_path, 'api_password: !secret pw')
         assert mock_error.call_count == 1, \
             "Expected an error about logger: value"
+
+
+def test_representing_yaml_loaded_data():
+    """Test we can represent YAML loaded data."""
+    files = {YAML_CONFIG_FILE: 'key: [1, "2", 3]'}
+    with patch_yaml_files(files):
+        data = load_yaml_config_file(YAML_CONFIG_FILE)
+    assert yaml.dump(data) == "key:\n- 1\n- '2'\n- 3\n"