feat: track original file in auto_yaml

mcbeet · Mar 1, 2022 · 36af60e · 36af60e
1 parent 54c46c4
commit 36af60e
Showing 1 changed file with 12 additions and 5 deletions.
diff --git a/beet/contrib/auto_yaml.py b/beet/contrib/auto_yaml.py
@@ -12,7 +12,6 @@
 from typing import Any, Tuple, Type
 
 from beet import Context, Drop, JsonFileBase, NamespaceFile, Pack, YamlFile
-from beet.core.utils import dump_json
 
 
 def beet_default(ctx: Context):
@@ -49,7 +48,7 @@ def use_auto_yaml(pack: Pack[Any]):
 
 
 def create_namespace_handler(
-    file_type: Type[NamespaceFile],
+    file_type: Type[JsonFileBase[Any]],
     namespace_scope: Tuple[str, ...],
     namespace_extension: str,
 ) -> Type[NamespaceFile]:
@@ -59,9 +58,11 @@ class AutoYamlNamespaceHandler(YamlFile, NamespaceFile):
         scope = namespace_scope
         extension = namespace_extension
 
+        model = file_type.model
+
         def bind(self, pack: Any, path: str):
             super().bind(pack, path)
-            pack[file_type].merge({path: file_type(dump_json(self.data))})
+            pack[file_type].merge({path: file_type(self.data, original=self.original)})
             raise Drop()
 
     return AutoYamlNamespaceHandler
@@ -74,9 +75,11 @@ def create_extra_handler(
     """Create handler that turns yaml extra files into json."""
 
     class AutoYamlExtraHandler(YamlFile):
+        model = file_type.model
+
         def bind(self, pack: Any, path: str):
             super().bind(pack, path)
-            pack.extra.merge({filename: file_type(dump_json(self.data))})
+            pack.extra.merge({filename: file_type(self.data, original=self.original)})
             raise Drop()
 
     return AutoYamlExtraHandler
@@ -89,10 +92,14 @@ def create_namespace_extra_handler(
     """Create handler that turns yaml namespace extra files into json."""
 
     class AutoYamlExtraHandler(YamlFile):
+        model = file_type.model
+
         def bind(self, pack: Any, path: str):
             super().bind(pack, path)
             namespace, _, path = path.partition(":")
-            pack[namespace].extra.merge({filename: file_type(dump_json(self.data))})
+            pack[namespace].extra.merge(
+                {filename: file_type(self.data, original=self.original)}
+            )
             raise Drop()
 
     return AutoYamlExtraHandler