Repository: smithy-lang/smithy
Pull request #1925 (2023-08-15 15:46:58)

Allow creating specific tree types
Adds method to TokenTree that allows creating a tree with a specific TreeType. Previously, there was no way to do this since the parse methods on TreeType are package private. This addition supports use cases where new trees need to be added to a root tree.

A test was added for this method, and other tests were updated to use it.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.



Patch:
diff --git a/smithy-syntax/src/main/java/software/amazon/smithy/syntax/TokenTree.java b/smithy-syntax/src/main/java/software/amazon/smithy/syntax/TokenTree.java
index 2967bde64..6eb84d00e 100644
--- a/smithy-syntax/src/main/java/software/amazon/smithy/syntax/TokenTree.java
+++ b/smithy-syntax/src/main/java/software/amazon/smithy/syntax/TokenTree.java
@@ -44,6 +44,27 @@ public interface TokenTree extends FromSourceLocation {
         return capturingTokenizer.getRoot();
     }
 
+    /**
+     * Create a TokenTree of the given {@link TreeType} from an {@link IdlTokenizer}.
+     *
+     * <p>The following example shows how to create a {@link TokenTree} for a trait.</p>
+     *
+     * <pre>{@code
+     * IdlTokenizer tokenizer = IdlTokenizer.create("@someTrait");
+     * TokenTree traitTree = TokenTree.of(tokenizer, TreeType.TRAIT);
+     * }</pre>
+     *
+     * @param tokenizer Tokenizer to traverse.
+     * @param type Type of tree to create.
+     * @return Returns the created tree.
+     */
+    static TokenTree of(IdlTokenizer tokenizer, TreeType type) {
+        CapturingTokenizer capturingTokenizer = new CapturingTokenizer(tokenizer);
+        type.parse(capturingTokenizer);
+        // The root of the tree is always IDL with children appended, so the first child is the one we want.
+        return capturingTokenizer.getRoot().getChildren().get(0);
+    }
+
     /**
      * Create a leaf tree from a single token.
      *


Test patch:
diff --git a/smithy-syntax/src/test/java/software/amazon/smithy/syntax/TokenTreeTest.java b/smithy-syntax/src/test/java/software/amazon/smithy/syntax/TokenTreeTest.java
index 569886180..d16ebed91 100644
--- a/smithy-syntax/src/test/java/software/amazon/smithy/syntax/TokenTreeTest.java
+++ b/smithy-syntax/src/test/java/software/amazon/smithy/syntax/TokenTreeTest.java
@@ -33,6 +33,18 @@ public class TokenTreeTest {
         assertThat(tree.getChildren().get(2).getType(), equalTo(TreeType.SHAPE_SECTION));
     }
 
+    @Test
+    public void createsFromTokenizerAndType() {
+        IdlTokenizer tokenizer = IdlTokenizer.create("@foo");
+        TokenTree tree = TokenTree.of(tokenizer, TreeType.TRAIT);
+
+        assertThat(tree.getType(), is(TreeType.TRAIT));
+        assertThat(tree.getChildren(), hasSize(2));
+        assertThat(tree.getError(), nullValue());
+        assertThat(tree.getChildren().get(0).getType(), equalTo(TreeType.TOKEN));
+        assertThat(tree.getChildren().get(1).getType(), equalTo(TreeType.SHAPE_ID));
+    }
+
     @Test
     public void createsFromCapturedToken() {
         IdlTokenizer tokenizer = IdlTokenizer.create("foo");
diff --git a/smithy-syntax/src/test/java/software/amazon/smithy/syntax/TreeTypeTest.java b/smithy-syntax/src/test/java/software/amazon/smithy/syntax/TreeTypeTest.java
index def758ddf..3f5f87640 100644
--- a/smithy-syntax/src/test/java/software/amazon/smithy/syntax/TreeTypeTest.java
+++ b/smithy-syntax/src/test/java/software/amazon/smithy/syntax/TreeTypeTest.java
@@ -1860,9 +1860,7 @@ public class TreeTypeTest {
     }
 
     private static TokenTree getTree(TreeType type, String forText) {
-        CapturingTokenizer tokenizer = new CapturingTokenizer(IdlTokenizer.create(forText));
-        type.parse(tokenizer);
-        // The root of the tree is always IDL with children appended, so the first child is the one we want.
-        return tokenizer.getRoot().getChildren().get(0);
+        IdlTokenizer tokenizer = IdlTokenizer.create(forText);
+        return TokenTree.of(tokenizer, type);
     }
 }