Avoid StopIteration when default YAML ends with comments

confuse/yaml_util.py

@@ -226,11 +226,13 @@ def restore_yaml_comments(data: str, default_data: str) -> str:
             comment = f"{line}\n"
         else:
             continue
-        while True:
-            line = next(default_lines)
+        for line in default_lines:
             if line and not line.startswith("#"):
                 break
             comment += f"{line}\n"
+        else:
+            # No following key: discard trailing comments and blank lines.
+            break
         key = line.split(":")[0].strip()
         comment_map[key] = comment
     out_lines = iter(data.splitlines())

docs/changelog.rst

@@ -13,9 +13,12 @@ Unreleased
     New features
     ~~~~~~~~~~~~
 
-..
-    Bug fixes
-    ~~~~~~~~~
+Bug fixes
+~~~~~~~~~
+
+- Fix ``Configuration.dump()`` raising ``StopIteration`` when the default
+  configuration ends with a comment or blank line (`#148
+  <https://github.com/beetbox/confuse/issues/148>`_).
 
 ..
     For plugin developers

test/test_dump.py

@@ -2,6 +2,9 @@
 import unittest
 from collections import OrderedDict
 
+import pytest
+import yaml
+
 import confuse
 
 from . import _root
@@ -106,3 +109,25 @@ def test_dump_redacted_sans_defaults(self):
 
         yaml = config.dump(redact=True, full=False).strip()
         assert yaml == "baz: REDACTED"
+
+
+@pytest.mark.parametrize(
+    "suffix", ["# trailing", "# trailing\n", "\n", "\n# trailing\n\n"]
+)
+def test_dump_defaults_with_trailing_comments(tmp_path, suffix):
+    path = tmp_path / "config_default.yaml"
+    path.write_text("# foo comment\nfoo: 42\n" + suffix, encoding="utf-8")
+    config = confuse.Configuration("myapp", read=False)
+    config.add(confuse.YamlSource(str(path), default=True))
+    config.set({"foo": 43})
+    assert config.dump() == "# foo comment\nfoo: 43\n"
+
+
+@pytest.mark.parametrize("style", ["|", ">"])
+def test_dump_defaults_preserves_block_scalar(tmp_path, style):
+    default_data = f"text: {style}\n  # scalar content\n  value\nother: 1\n"
+    path = tmp_path / "config_default.yaml"
+    path.write_text(default_data, encoding="utf-8")
+    config = confuse.Configuration("myapp", read=False)
+    config.add(confuse.YamlSource(str(path), default=True))
+    assert yaml.safe_load(config.dump()) == yaml.safe_load(default_data)

test/test_yaml.py

@@ -114,3 +114,39 @@ def test_list_unchanged(self):
     def test_invalid_yaml_string_unchanged(self):
         v = confuse.yaml_util.parse_as_scalar("!", confuse.Loader)
         assert v == "!"
+
+
+@pytest.mark.parametrize(
+    "suffix", ["# trailing", "# trailing\n", "\n", "\n\n", "\n# trailing\n\n"]
+)
+def test_restore_comments_at_eof(suffix):
+    default_data = "# before foo\nfoo: 42\n" + suffix
+    data = "foo: 43\n"
+    assert confuse.yaml_util.restore_yaml_comments(data, default_data) == (
+        "# before foo\nfoo: 43\n"
+    )
+
+
+@pytest.mark.parametrize(
+    "default_data", ["", "\n", "# comment", "# comment\n\n# another\n"]
+)
+def test_restore_comments_without_keys(default_data):
+    data = "foo: 42\n"
+    assert confuse.yaml_util.restore_yaml_comments(data, default_data) == data
+
+
+@pytest.mark.parametrize("style", ["|", ">"])
+def test_restore_comments_preserves_block_scalar(style):
+    default_data = f"text: {style}\n  # scalar content\n  value\nother: 1\n"
+    expected = load(default_data)
+    data = yaml.dump(expected, Dumper=confuse.Dumper)
+    restored = confuse.yaml_util.restore_yaml_comments(data, default_data)
+    assert load(restored) == expected
+
+
+def test_restore_comments_before_reordered_keys():
+    default_data = "# foo comment\nfoo: 42\n\n# bar comment\nbar: 1\n"
+    data = "bar: 2\nfoo: 43\n"
+    assert confuse.yaml_util.restore_yaml_comments(data, default_data) == (
+        "\n# bar comment\nbar: 2\n# foo comment\nfoo: 43\n"
+    )