Repository: kokorin/dbt-pumpkin

Issue:
DiskStorage: use round-trip YAML
It appears ruamel provides dedicated parser/dumper to read and update YAML files. It also provides limited capabilities to keep comments (at least line-ending).

`YAML(typ="rt")` or just `YAML()`

Discussion:

Patch:
diff --git a/README.md b/README.md
index 3eb41d9..4bfa941 100644
--- a/README.md
+++ b/README.md
@@ -8,7 +8,7 @@ Inspired by [dbt-osmosis](https://z3z1ma.github.io/dbt-osmosis/)
 
 ### Botstrap DBT Resources
 
-`dbt-osmosis` allows to create DBT YAML schema files for Seeds, Models and Snapshots. For that one has to add
+`dbt-pumpkin` allows to create DBT YAML schema files for Seeds, Models and Snapshots. For that one has to add
 [dbt-pumpking-path](#dbt-pumpkin-path) configuration property for a Resource. Then `dbt-pumpkin` will be able to
 understand where YAML files should be located.
 
@@ -39,7 +39,7 @@ Options:
 
 ### Relocate DBT Resources
 
-`dbt-osmosis` also allows to move DBT YAML schema files for Sources, Seeds, Models and Snapshots. As
+`dbt-pumpkin` also allows to move DBT YAML schema files for Sources, Seeds, Models and Snapshots. As
 with [bootstrap](#botstrap-dbt-resources) command [dbt-pumpking-path](#dbt-pumpkin-path) configuration property is
 required for a Resource. Then `dbt-pumpkin` will be able to understand where YAML files should be located.
 
diff --git a/dbt_pumpkin/storage.py b/dbt_pumpkin/storage.py
index b48a759..765ce43 100644
--- a/dbt_pumpkin/storage.py
+++ b/dbt_pumpkin/storage.py
@@ -22,7 +22,7 @@ class DiskStorage(Storage):
     def __init__(self, root_dir: Path, *, read_only: bool):
         self._root_dir = root_dir
         self._read_only = read_only
-        self._yaml = YAML(typ="safe")
+        self._yaml = YAML(typ="rt")
 
     def load_yaml(self, files: set[Path]) -> dict[Path, Any]:
         result: dict[Path, Any] = {}


Test patch:
diff --git a/tests/test_storage.py b/tests/test_storage.py
index ec1b1dd..c558009 100644
--- a/tests/test_storage.py
+++ b/tests/test_storage.py
@@ -1,3 +1,4 @@
+import platform
 import textwrap
 from pathlib import Path
 
@@ -6,7 +7,7 @@ import yaml
 from dbt_pumpkin.storage import DiskStorage
 
 
-def test_load_yaml(tmp_path):
+def test_load_yaml(tmp_path: Path):
     (tmp_path / "schema.yml").write_text(
         textwrap.dedent("""\
         version: 2
@@ -19,7 +20,7 @@ def test_load_yaml(tmp_path):
     assert files == {Path("schema.yml"): {"version": 2, "models": [{"name": "my_model"}]}}
 
 
-def test_save_yaml(tmp_path):
+def test_save_yaml(tmp_path: Path):
     DiskStorage(tmp_path, read_only=False).save_yaml(
         {Path("schema.yml"): {"version": 2, "models": [{"name": "my_other_model"}]}}
     )
@@ -28,9 +29,52 @@ def test_save_yaml(tmp_path):
     assert actual == {"version": 2, "models": [{"name": "my_other_model"}]}
 
 
-def test_save_yaml_read_only(tmp_path):
+def test_save_yaml_read_only(tmp_path: Path):
     DiskStorage(tmp_path, read_only=True).save_yaml(
         {Path("schema.yml"): {"version": 2, "models": [{"name": "my_other_model"}]}}
     )
 
     assert not (tmp_path / "schema.yml").exists()
+
+
+def test_roundtrip(tmp_path: Path):
+    content = textwrap.dedent("""\
+        version: 2
+        models:
+        # TODO rename it!
+        - name: my_model
+          description: my very first model
+          columns:
+          - name: id
+            data_type: short # or is it int actually?
+          - name: name
+    """)
+
+    (tmp_path / "my_model.yml").write_text(content)
+
+    storage = DiskStorage(tmp_path, read_only=False)
+    files = storage.load_yaml({Path("my_model.yml")})
+    storage.save_yaml(files)
+
+    actual = (tmp_path / "my_model.yml").read_text()
+
+    # Ruamel adds newlines after comments on Windows, but not on Ubuntu
+    # It should be fine as we don't re-write a file if there is nothing to change in it
+    if platform.system().lower() == "windows":
+        expected = textwrap.dedent("""\
+            version: 2
+            models:
+            # TODO rename it!
+
+            - name: my_model
+              description: my very first model
+              columns:
+              - name: id
+                data_type: short # or is it int actually?
+
+              - name: name
+        """)
+    else:
+        expected = content
+
+    assert expected == actual