Fix StopIteration in restore_yaml_comments when default_data ends with a comment or blank line

confuse/yaml_util.py

@@ -227,10 +227,18 @@ def restore_yaml_comments(data: str, default_data: str) -> str:
         else:
             continue
         while True:
-            line = next(default_lines)
+            try:
+                line = next(default_lines)
+            except StopIteration:
+                # File ends with a comment or blank line and no following key;
+                # discard the accumulated comment and stop scanning.
+                line = None
+                break
             if line and not line.startswith("#"):
                 break
             comment += f"{line}\n"
+        if line is None:
+            break
         key = line.split(":")[0].strip()
         comment_map[key] = comment
     out_lines = iter(data.splitlines())

docs/changelog.rst

@@ -6,6 +6,8 @@ Unreleased
 
 - Require `typing_extensions` on older Python versions (and use `typing` on
   newer versions). [#189](https://github.com/beetbox/confuse/issues/189)
+- Fix ``restore_yaml_comments`` crashing when default data ends with a comment
+  or blank line. [#148](https://github.com/beetbox/confuse/issues/148)
 
 v2.2.0
 ------

test/test_yaml.py

@@ -5,6 +5,7 @@
 import yaml
 
 import confuse
+from confuse.yaml_util import restore_yaml_comments
 
 from . import TempDir
 
@@ -114,3 +115,30 @@ def test_list_unchanged(self):
     def test_invalid_yaml_string_unchanged(self):
         v = confuse.yaml_util.parse_as_scalar("!", confuse.Loader)
         assert v == "!"
+
+
+class RestoreYamlCommentsTest(unittest.TestCase):
+    def test_trailing_comment_does_not_raise(self):
+        """restore_yaml_comments must not raise StopIteration when
+        default_data ends with a comment line (issue #148)."""
+        default_data = "key1: value1\n# trailing comment\n"
+        data = "key1: value1\n"
+        # Must complete without raising StopIteration
+        result = restore_yaml_comments(data, default_data)
+        assert "key1: value1" in result
+
+    def test_trailing_blank_line_does_not_raise(self):
+        """restore_yaml_comments must not raise StopIteration when
+        default_data ends with a blank line (issue #148)."""
+        default_data = "key1: value1\n\n"
+        data = "key1: value1\n"
+        result = restore_yaml_comments(data, default_data)
+        assert "key1: value1" in result
+
+    def test_comments_still_restored_before_key(self):
+        """Comments before a key in default_data are prepended to that key
+        in the output."""
+        default_data = "# Comment for key1\nkey1: value1\n"
+        data = "key1: value1\n"
+        result = restore_yaml_comments(data, default_data)
+        assert result == "# Comment for key1\nkey1: value1\n"