add custom yaml tags initial support

Author: mr-mixas

nested_diff/cli.py

@@ -19,6 +19,7 @@
 import sys
 
 import nested_diff
+import nested_diff.handlers
 
 
 class App:
@@ -220,6 +221,7 @@ def load(self, file_):
             fmt = self.args.ifmt
 
         fmt_opts = self._decode_fmt_opts(self.args.ifmt_opts)
+
         return self.get_loader(fmt, **fmt_opts).load(file_)
 
     @staticmethod
@@ -550,15 +552,105 @@ def __init__(self, **kwargs):
 
         import yaml
 
+        from yaml.nodes import (
+            ScalarNode as YamlScalarNode,
+            SequenceNode as YamlSequenceNode,
+            MappingNode as YamlMappingNode,
+        )
+
         try:
             from yaml import CSafeLoader as YamlLoader
         except ImportError:
             from yaml import SafeLoader as YamlLoader
 
+        self.opts = self.get_opts(kwargs)
         self.yaml = yaml
         self.yaml_loader = YamlLoader
-        self.opts = self.get_opts(kwargs)
+
+        def _default_constructor(loader, tag_suffix, node):  # noqa: ARG001
+            tag = node.tag
+
+            if isinstance(node, YamlScalarNode):
+                value = node.value
+            elif isinstance(node, YamlSequenceNode):
+                node.tag = 'tag:yaml.org,2002:seq'
+                value = loader.construct_sequence(node, deep=True)
+            elif isinstance(node, YamlMappingNode):
+                node.tag = 'tag:yaml.org,2002:map'
+                value = loader.construct_mapping(node, deep=True)
+
+            return YamlNode(tag, value)
+
+        self.yaml_loader.add_multi_constructor(None, _default_constructor)
 
     def decode(self, data):
         """Parse YAML string."""
         return self.yaml.load(data, Loader=self.yaml_loader, **self.opts)
+
+
+class YamlNode:
+    """Wrapper to represent YAML node."""
+
+    def __init__(self, tag, value):
+        """Initialize wrapper.
+
+        Args:
+            tag: YAML node tag.
+            value: YAML node value.
+
+        """
+        self.tag = tag
+        self.value = value
+
+    def __repr__(self):
+        """Repr for YAML node wrapper."""
+        return f"YamlNode(tag='{self.tag}', value={self.value!r})"
+
+
+class YamlNodeHandler(nested_diff.handlers.TypeHandler):
+    """YamlNode handler."""
+
+    extension_id = 'nested_diff.YamlNode'
+    handled_type = YamlNode
+
+    def diff(self, differ, a, b):
+        """Calculate diff for two YamlNode objects.
+
+        Args:
+            differ: nested_diff.Differ object.
+            a: First node to diff.
+            b: Second node to diff.
+
+        Returns:
+            Tuple: equality flag and nested diff.
+
+        """
+        equal, _ = differ.diff(a.tag, b.tag)
+
+        if not equal:
+            return equal, {'N': b, 'O': a, 'E': self.extension_id}
+
+        equal, diff = differ.diff(a.value, b.value)
+
+        if diff:
+            diff = {
+                'D': {'value': diff},
+                'E': self.extension_id,
+                'tag': a.tag,
+            }
+
+        return equal, diff
+
+    def generate_formatted_diff(self, formatter, diff, depth):
+        """Generate formatted YamlNode diff."""
+        if 'D' in diff:
+            yield from formatter.generate_string(diff['tag'], 'D', depth)
+            yield from formatter.generate_diff(diff['D']['value'], depth + 1)
+
+            return
+
+        yield from formatter.generate_string(diff['O'].tag, 'O', depth)
+        yield from formatter.generate_value(diff['O'].value, 'O', depth + 1)
+
+        yield from formatter.generate_string(diff['N'].tag, 'N', depth)
+        yield from formatter.generate_value(diff['N'].value, 'N', depth + 1)

nested_diff/diff_tool.py

@@ -57,6 +57,7 @@ def diff(self, a, b, **kwargs):
             diff_opts['R'] = int(diff_opts['R'])
 
         differ = nested_diff.Differ(**diff_opts)
+        differ.set_handler(nested_diff.cli.YamlNodeHandler())
 
         if self.args.text_ctx >= 0:
             differ.set_handler(
@@ -286,6 +287,7 @@ def __init__(
 
         fmt_class = self.get_formatter_class(base_class, values=values)
         self.encoder = fmt_class(**self.get_opts(kwargs))
+        self.encoder.set_handler(nested_diff.cli.YamlNodeHandler())
 
         if header is None:
             if hasattr(self.encoder, 'get_page_header'):

pyproject.toml

@@ -44,7 +44,7 @@ requires-python = '>=3.7'
 
 [project.optional-dependencies]
 cli = [
-    'pyyaml',
+    'pyyaml >= 5.3',
     'tomli >= 1.1.0 ; python_version < "3.11"',
     'tomli-w >= 1.0.0'
 ]

tests/cli/shared.custom_tags.a.yaml

@@ -0,0 +1,15 @@
+---
+changed_value: !changed_value old_value
+changed_tag: !changed_tag changed_tag_value
+unchanged_value: !unchanged_value unchanged_value
+removed: !removed
+
+recursive_mapping: !mapping
+  child: !mapping
+    var: bar
+  var: foo
+
+sequence: !sequence
+  - foo
+  - bar
+  - baz

tests/cli/shared.custom_tags.b.yaml

@@ -0,0 +1,15 @@
+---
+added: !added
+changed_value: !changed_value new_value
+changed_tag: !Changed_tag changed_tag_value
+unchanged_value: !unchanged_value unchanged_value
+
+recursive_mapping: !mapping
+  child: !mapping
+    var: baz
+  var: foo
+
+sequence: !sequence
+  - foo
+  - baz
+  - bar

tests/cli/test_cli.py

@@ -110,6 +110,14 @@ def test_get_loader_unsupported_fmt():
         cli.App(args=()).get_loader('garbage')
 
 
+def test_loader_yaml_custom_tags():
+    loaded = cli.YamlLoader().decode('!custom_tag')
+
+    assert isinstance(loaded, cli.YamlNode)
+    assert loaded.tag == '!custom_tag'
+    assert loaded.value == ''
+
+
 def test_run():
     with pytest.raises(NotImplementedError):
         cli.App(args=()).run()

tests/cli/test_diff_tool.py

@@ -748,3 +748,20 @@ def test_values_yaml_multiline_strings(capsys, expected, rpath):
     assert exit_code == 0
 
     assert captured.out == expected
+
+
+def test_yaml_custom_tags(capsys, expected, rpath):
+    exit_code = nested_diff.diff_tool.App(
+        args=(
+            rpath('shared.custom_tags.a.yaml'),
+            rpath('shared.custom_tags.b.yaml'),
+            '-U',
+            '1',
+        ),
+    ).run()
+
+    captured = capsys.readouterr()
+    assert captured.err == ''
+    assert exit_code == 1
+
+    assert captured.out == expected

tests/cli/test_diff_tool.test_yaml_custom_tags.exp

@@ -0,0 +1,41 @@
++ {'added'}
++   YamlNode(tag='!added', value='')
+  {'changed_tag'}
+#   <YamlNode>
+-   !changed_tag
+-     'changed_tag_value'
++   !Changed_tag
++     'changed_tag_value'
+  {'changed_value'}
+#   <YamlNode>
+    !changed_value
+-     'old_value'
++     'new_value'
+  {'recursive_mapping'}
+#   <YamlNode>
+    !mapping
+      {'child'}
+#       <YamlNode>
+        !mapping
+          {'var'}
+-           'bar'
++           'baz'
+      {'var'}
+        'foo'
+- {'removed'}
+-   YamlNode(tag='!removed', value='')
+  {'sequence'}
+#   <YamlNode>
+    !sequence
+      [0]
+        'foo'
++     [1]
++       'baz'
+      [2]
+        'bar'
+-     [3]
+-       'baz'
+  {'unchanged_value'}
+#   <YamlNode>
+    !unchanged_value
+      'unchanged_value'