Amending YAML parser to add support for tags

rewrite-yaml/src/main/java/org/openrewrite/yaml/AppendToSequenceVisitor.java

@@ -112,7 +112,7 @@ private Yaml.Sequence appendToSequence(Yaml.Sequence existingSequence, String va
                 entries.set(lastEntryIndex, existingEntry.withTrailingCommaPrefix(""));
             }
         }
-        Yaml.Scalar newItem = new Yaml.Scalar(randomId(), itemPrefix, Markers.EMPTY, style, null, value);
+        Yaml.Scalar newItem = new Yaml.Scalar(randomId(), itemPrefix, Markers.EMPTY, style, null, null, value);
         Yaml.Sequence.Entry newEntry = new Yaml.Sequence.Entry(randomId(), entryPrefix, Markers.EMPTY, newItem, hasDash, entryTrailingCommaPrefix);
         entries.add(newEntry);
         return newSequence.withMarkers(Markers.EMPTY.addIfAbsent(new AlreadyReplaced(randomId(), value, value)));

rewrite-yaml/src/main/java/org/openrewrite/yaml/ChangePropertyKey.java

@@ -251,7 +251,7 @@ public Yaml.Mapping visitMapping(Yaml.Mapping mapping, P p) {
                         newEntryPrefix,
                         Markers.EMPTY,
                         new Yaml.Scalar(randomId(), "", Markers.EMPTY,
-                                Yaml.Scalar.Style.PLAIN, null, subproperty),
+                                Yaml.Scalar.Style.PLAIN, null, null, subproperty),
                         scope.getBeforeMappingValueIndicator(),
                         removeExclusions(entryToReplace.getValue().copyPaste()));
 

rewrite-yaml/src/main/java/org/openrewrite/yaml/ChangeValue.java

@@ -69,10 +69,11 @@ public Yaml.Mapping.Entry visitMappingEntry(Yaml.Mapping.Entry entry, ExecutionC
                 Yaml.Mapping.Entry e = super.visitMappingEntry(entry, ctx);
                 if (matcher.matches(getCursor()) && (!(e.getValue() instanceof Yaml.Scalar) || !((Yaml.Scalar) e.getValue()).getValue().equals(value))) {
                     Yaml.Anchor anchor = (e.getValue() instanceof Yaml.Scalar) ? ((Yaml.Scalar) e.getValue()).getAnchor() : null;
+                    Yaml.Tag tag = (e.getValue() instanceof Yaml.Scalar) ? ((Yaml.Scalar) e.getValue()).getTag() : null;
                     String prefix = e.getValue() instanceof Yaml.Sequence ? ((Yaml.Sequence) e.getValue()).getOpeningBracketPrefix() : e.getValue().getPrefix();
                     e = e.withValue(
                             new Yaml.Scalar(randomId(), prefix, Markers.EMPTY,
-                                    Yaml.Scalar.Style.PLAIN, anchor, value)
+                                    Yaml.Scalar.Style.PLAIN, anchor, tag, value)
                     );
                 }
                 return e;

rewrite-yaml/src/main/java/org/openrewrite/yaml/CountLinesVisitor.java

@@ -85,6 +85,15 @@ public Yaml visitAnchor(Yaml.Anchor anchor, AtomicInteger count) {
         return super.visitAnchor(anchor, count);
     }
 
+    @Override
+    public Yaml visitTag(Yaml.Tag tag, AtomicInteger count) {
+        if(tag.getPrefix().contains("\n")) {
+            count.incrementAndGet();
+        }
+        return super.visitTag(tag, count);
+    }
+
+
     @Override
     public Yaml visitAlias(Yaml.Alias alias, AtomicInteger count) {
         if(alias.getPrefix().contains("\n")) {

rewrite-yaml/src/main/java/org/openrewrite/yaml/YamlIsoVisitor.java

@@ -57,4 +57,9 @@ public Yaml.Sequence visitSequence(Yaml.Sequence sequence, P p) {
     public Yaml.Sequence.Entry visitSequenceEntry(Yaml.Sequence.Entry entry, P p) {
         return (Yaml.Sequence.Entry) super.visitSequenceEntry(entry, p);
     }
+
+    @Override
+    public Yaml.Tag visitTag(Yaml.Tag tag, P p) {
+        return (Yaml.Tag) super.visitTag(tag, p);
+    }
 }

rewrite-yaml/src/main/java/org/openrewrite/yaml/YamlParser.java

@@ -209,6 +209,25 @@ private Yaml.Documents parseFromInput(Path sourceFile, EncodingDetectingInputStr
                         valueStart = valueStart - newLine.length();
                         newLine = "";
 
+                        Yaml.Tag tag = null;
+                        if (scalar.getTag() != null) {
+                            String potentialScalarValue = reader.readStringFromBuffer(valueStart, event.getEndMark().getIndex() - 1);
+                            assert(potentialScalarValue.contains("!"));
+                            final int tagStartIndex = potentialScalarValue.indexOf('!');
+                            String tagPrefix = potentialScalarValue.substring(0, tagStartIndex);
+                            int indexOfTagName = tagStartIndex;
+                            while (indexOfTagName < potentialScalarValue.length() && !Character.isWhitespace(potentialScalarValue.charAt(indexOfTagName))) {
+                                indexOfTagName++;
+                            }
+                            String tagName = potentialScalarValue.substring(tagStartIndex, indexOfTagName);
+                            int indexOfSpaceAfterTag = indexOfTagName;
+                            while (indexOfSpaceAfterTag < potentialScalarValue.length() && Character.isWhitespace(potentialScalarValue.charAt(indexOfSpaceAfterTag))) {
+                                indexOfSpaceAfterTag++;
+                            }
+                            String tagSuffix = potentialScalarValue.substring(indexOfTagName, indexOfSpaceAfterTag);
+                            valueStart = valueStart + indexOfSpaceAfterTag;
+                            tag = new Yaml.Tag(randomId(), tagPrefix, tagSuffix, Markers.EMPTY, tagName);
+                        }
 
                         String scalarValue;
                         switch (scalar.getScalarStyle()) {
@@ -252,12 +271,9 @@ private Yaml.Documents parseFromInput(Path sourceFile, EncodingDetectingInputStr
                             case PLAIN:
                             default:
                                 style = Yaml.Scalar.Style.PLAIN;
-                                if (!scalarValue.startsWith("@") && event.getStartMark().getIndex() >= reader.getBufferIndex()) {
-                                    scalarValue = reader.readStringFromBuffer(event.getStartMark().getIndex(), event.getEndMark().getIndex() - 1);
-                                }
                                 break;
                         }
-                        Yaml.Scalar finalScalar = new Yaml.Scalar(randomId(), fmt, Markers.EMPTY, style, anchor, scalarValue);
+                        Yaml.Scalar finalScalar = new Yaml.Scalar(randomId(), fmt, Markers.EMPTY, style, anchor, tag, scalarValue);
                         BlockBuilder builder = blockStack.isEmpty() ? null : blockStack.peek();
                         if (builder instanceof SequenceBuilder) {
                             // Inline sequences like [1, 2] need to keep track of any whitespace between the element
@@ -329,19 +345,32 @@ private Yaml.Documents parseFromInput(Path sourceFile, EncodingDetectingInputStr
                                 fmt = fmt.substring(0, dashPrefixIndex);
                             }
                         }
-
                         String fullPrefix = reader.readStringFromBuffer(lastEnd, event.getEndMark().getIndex());
                         String startBracketPrefix = null;
                         int openingBracketIndex = commentAwareIndexOf('[', fullPrefix);
+                        int startIndex = commentAwareIndexOf(':', fullPrefix) + 1;
                         if (openingBracketIndex != -1) {
-                            int startIndex = commentAwareIndexOf(':', fullPrefix) + 1;
                             startBracketPrefix = fullPrefix.substring(startIndex, openingBracketIndex);
                         }
+                        Yaml.Tag tag = null;
+                        if (sse.getTag() != null) {
+                            String prefixAfterColon = fullPrefix.substring(startIndex);
+                            final int tagStartIndex = prefixAfterColon.indexOf('!');
+                            String tagPrefix = prefixAfterColon.substring(0, tagStartIndex);
+                            int i = tagStartIndex;
+                            while (i < prefixAfterColon.length() && !Character.isWhitespace(prefixAfterColon.charAt(i))) {
+                                i++;
+                            }
+                            // Cannot use sse.getTag() here, because it is sometimes expanded, e.g. `!!seq` becomes `tag:yaml.org,2002:seq`
+                            String tagName = prefixAfterColon.substring(tagStartIndex, i);
+                            String tagSuffix = prefixAfterColon.substring(i, prefixAfterColon.length() - 2);
+                            tag = new Yaml.Tag(randomId(), tagPrefix, tagSuffix, Markers.EMPTY, tagName);
+                        }
                         lastEnd = event.getEndMark().getIndex();
                         if (shouldUseYamlParserBugWorkaround(sse)) {
                             lastEnd--;
                         }
-                        blockStack.push(new SequenceBuilder(fmt, startBracketPrefix, anchor));
+                        blockStack.push(new SequenceBuilder(fmt, startBracketPrefix, anchor, tag));
                         break;
                     }
                     case Alias: {
@@ -512,13 +541,15 @@ private static class SequenceBuilder implements BlockBuilder {
 
 
         private final Yaml.@Nullable Anchor anchor;
+        private final Yaml.@Nullable Tag tag;
 
         private final List<Yaml.Sequence.Entry> entries = new ArrayList<>();
 
-        private SequenceBuilder(String prefix, @Nullable String startBracketPrefix, Yaml.@Nullable Anchor anchor) {
+        private SequenceBuilder(String prefix, @Nullable String startBracketPrefix, Yaml.@Nullable Anchor anchor, Yaml.@Nullable Tag tag) {
             this.prefix = prefix;
             this.startBracketPrefix = startBracketPrefix;
             this.anchor = anchor;
+            this.tag = tag;
         }
 
         @Override
@@ -544,7 +575,7 @@ public void push(Yaml.Block block, @Nullable String commaPrefix) {
 
         @Override
         public SequenceWithPrefix build() {
-            return new SequenceWithPrefix(prefix, startBracketPrefix, entries, null, anchor);
+            return new SequenceWithPrefix(prefix, startBracketPrefix, entries, null, anchor, tag);
         }
     }
 
@@ -584,8 +615,8 @@ public Mapping withPrefix(String prefix) {
     private static class SequenceWithPrefix extends Yaml.Sequence {
         private String prefix;
 
-        public SequenceWithPrefix(String prefix, @Nullable String startBracketPrefix, List<Yaml.Sequence.Entry> entries, @Nullable String endBracketPrefix, @Nullable Anchor anchor) {
-            super(randomId(), Markers.EMPTY, startBracketPrefix, entries, endBracketPrefix, anchor);
+        public SequenceWithPrefix(String prefix, @Nullable String startBracketPrefix, List<Yaml.Sequence.Entry> entries, @Nullable String endBracketPrefix, @Nullable Anchor anchor, @Nullable Tag tag) {
+            super(randomId(), Markers.EMPTY, startBracketPrefix, entries, endBracketPrefix, anchor, tag);
             this.prefix = prefix;
         }
 
@@ -610,7 +641,8 @@ public Yaml.Sequence visitSequence(Yaml.Sequence sequence, Integer p) {
                                     sequenceWithPrefix.getOpeningBracketPrefix(),
                                     ListUtils.mapFirst(sequenceWithPrefix.getEntries(), e -> e.withPrefix(sequenceWithPrefix.getPrefix())),
                                     sequenceWithPrefix.getClosingBracketPrefix(),
-                                    sequenceWithPrefix.getAnchor()
+                                    sequenceWithPrefix.getAnchor(),
+                                    sequenceWithPrefix.getTag()
                             ), p);
                 }
                 return super.visitSequence(sequence, p);

rewrite-yaml/src/main/java/org/openrewrite/yaml/YamlVisitor.java

@@ -92,12 +92,16 @@ public Yaml visitMappingEntry(Yaml.Mapping.Entry entry, P p) {
     }
 
     public Yaml visitScalar(Yaml.Scalar scalar, P p) {
-        return scalar.withAnchor(visitAndCast(scalar.getAnchor(), p))
+        return scalar
+                .withAnchor(visitAndCast(scalar.getAnchor(), p))
+                .withTag(visitAndCast(scalar.getTag(), p))
                 .withMarkers(visitMarkers(scalar.getMarkers(), p));
     }
 
     public Yaml visitSequence(Yaml.Sequence sequence, P p) {
-        return sequence.withEntries(ListUtils.map(sequence.getEntries(), e -> visitAndCast(e, p)))
+        return sequence
+                .withTag(visitAndCast(sequence.getTag(), p))
+                .withEntries(ListUtils.map(sequence.getEntries(), e -> visitAndCast(e, p)))
                 .withMarkers(visitMarkers(sequence.getMarkers(), p));
     }
 
@@ -115,6 +119,10 @@ public Yaml visitAlias(Yaml.Alias alias, P p) {
                 .withMarkers(visitMarkers(alias.getMarkers(), p));
     }
 
+    public Yaml visitTag(Yaml.Tag tag, P p) {
+        return tag.withMarkers(visitMarkers(tag.getMarkers(), p));
+    }
+
     @Deprecated
     public void maybeCoalesceProperties() {
         if (getAfterVisit().stream().noneMatch(CoalescePropertiesVisitor.class::isInstance)) {

rewrite-yaml/src/main/java/org/openrewrite/yaml/internal/YamlPrinter.java

@@ -116,6 +116,9 @@ public Yaml visitMapping(Yaml.Mapping mapping, PrintOutputCapture<P> p) {
     @Override
     public Yaml visitScalar(Yaml.Scalar scalar, PrintOutputCapture<P> p) {
         beforeSyntax(scalar, p);
+        if (scalar.getTag() != null) {
+            visit(scalar.getTag(), p);
+        }
         if (scalar.getAnchor() != null) {
             visit(scalar.getAnchor(), p);
         }
@@ -170,6 +173,17 @@ public Yaml visitAlias(Yaml.Alias alias, PrintOutputCapture<P> p) {
         return alias;
     }
 
+    @Override
+    public Yaml visitTag(Yaml.Tag tag, PrintOutputCapture<P> p) {
+        visitMarkers(tag.getMarkers(), p);
+        p
+                .append(tag.getPrefix())
+                .append(tag.getName())
+                .append(tag.getSuffix());
+        afterSyntax(tag, p);
+        return tag;
+    }
+
     private static final UnaryOperator<String> YAML_MARKER_WRAPPER =
             out -> "~~" + out + (out.isEmpty() ? "" : "~~") + ">";
 

rewrite-yaml/src/main/java/org/openrewrite/yaml/tree/Yaml.java

@@ -230,6 +230,9 @@ class Scalar implements Block, YamlKey {
         @Nullable
         Anchor anchor;
 
+        @Nullable
+        Tag tag;
+
         String value;
 
         public enum Style {
@@ -247,7 +250,7 @@ public <P> Yaml acceptYaml(YamlVisitor<P> v, P p) {
 
         @Override
         public Scalar copyPaste() {
-            return new Scalar(randomId(), prefix, Markers.EMPTY, style, anchor, value);
+            return new Scalar(randomId(), prefix, Markers.EMPTY, style, anchor, tag, value);
         }
 
         @Override
@@ -362,6 +365,9 @@ class Sequence implements Block {
         @Nullable
         Anchor anchor;
 
+        @Nullable
+        Tag tag;
+
         @Override
         public <P> Yaml acceptYaml(YamlVisitor<P> v, P p) {
             return v.visitSequence(this, p);
@@ -370,7 +376,7 @@ public <P> Yaml acceptYaml(YamlVisitor<P> v, P p) {
         @Override
         public Sequence copyPaste() {
             return new Sequence(randomId(), Markers.EMPTY, openingBracketPrefix,
-                    entries.stream().map(Entry::copyPaste).collect(toList()), closingBracketPrefix, anchor);
+                    entries.stream().map(Entry::copyPaste).collect(toList()), closingBracketPrefix, anchor, tag);
         }
 
         @Override
@@ -506,4 +512,39 @@ public String toString() {
             return "Yaml.Anchor(" + key + ")";
         }
     }
+
+    @Value
+    @EqualsAndHashCode(callSuper = false, onlyExplicitlyIncluded = true)
+    @With
+    class Tag implements Yaml {
+        @EqualsAndHashCode.Include
+        UUID id;
+
+        @With
+        String prefix;
+
+        @With
+        String suffix;
+
+        @With
+        Markers markers;
+
+        @With
+        String name;
+
+        @Override
+        public <P> Yaml acceptYaml(YamlVisitor<P> v, P p) {
+            return v.visitTag(this, p);
+        }
+
+        @Override
+        public Tag copyPaste() {
+            return new Tag(randomId(), prefix, suffix, Markers.EMPTY, name);
+        }
+
+        @Override
+        public String toString() {
+            return "Yaml.Tag(" + name + ")";
+        }
+    }
 }

rewrite-yaml/src/test/java/org/openrewrite/yaml/YamlParserTest.java

@@ -237,12 +237,55 @@ void atSymbols() {
     void pipeLiteralInASequenceWithDoubleQuotes() {
         rewriteRun(
           yaml(
-            """
+               """
                - "one": |
                    two
                  "three": "four"
                """
           )
         );
     }
+
+    @Test
+    void tagsAsInCloudFormation() {
+        rewriteRun(
+          yaml(
+            """
+            AttributeDefinitions: !Dynamo
+              - AttributeName: Title
+            """
+          )
+        );
+    }
+
+    @Test
+    void tagsAsInScalar() {
+        rewriteRun(
+          yaml(
+            """
+            AttributeDefinitions: !Dynamo Title
+            """
+          )
+        );
+    }
+
+    @Test
+    void globalTags() {
+        rewriteRun(
+          yaml(
+            """
+            age: !!int "42"
+            pi: !!float "3.14159"
+            is_valid: !!bool "true"
+            names: !!seq
+              - Alice
+              - Bob
+              - Charlie
+            person: !!map
+              name: John Doe
+              age: 30
+            """
+          )
+        );
+    }
 }