Generate Union Documentation and Output Usage Example

codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/CodegenVisitor.java

@@ -251,7 +251,9 @@ public Void stringShape(StringShape shape) {
 
     @Override
     public Void unionShape(UnionShape shape) {
-        writers.useShapeWriter(shape, writer -> new UnionGenerator(model, symbolProvider, writer, shape).run());
+        UnionGenerator generator = new UnionGenerator(model, symbolProvider, shape);
+        writers.useShapeWriter(shape, generator::generateUnion);
+        writers.useShapeExportedTestWriter(shape, generator::generateUnionExamples);
         return null;
     }
 

codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/GoDelegator.java

@@ -70,7 +70,7 @@ List<SymbolDependency> getDependencies() {
     /**
      * Gets a previously created writer or creates a new one if needed.
      *
-     * @param shape Shape to create the writer for.
+     * @param shape          Shape to create the writer for.
      * @param writerConsumer Consumer that accepts and works with the file.
      */
     public void useShapeWriter(Shape shape, Consumer<GoWriter> writerConsumer) {
@@ -81,7 +81,7 @@ public void useShapeWriter(Shape shape, Consumer<GoWriter> writerConsumer) {
     /**
      * Gets a previously created writer or creates a new one for the a Go test file for the associated shape.
      *
-     * @param shape Shape to create the writer for.
+     * @param shape          Shape to create the writer for.
      * @param writerConsumer Consumer that accepts and works with the file.
      */
     public void useShapeTestWriter(Shape shape, Consumer<GoWriter> writerConsumer) {
@@ -93,16 +93,39 @@ public void useShapeTestWriter(Shape shape, Consumer<GoWriter> writerConsumer) {
         filename = b.toString();
 
         symbol = symbol.toBuilder()
-            .definitionFile(filename)
-        .build();
+                .definitionFile(filename)
+                .build();
+
+        useShapeWriter(symbol, writerConsumer);
+    }
+
+    /**
+     * Gets a previously created writer or creates a new one for the a Go public package test file for the associated
+     * shape.
+     *
+     * @param shape          Shape to create the writer for.
+     * @param writerConsumer Consumer that accepts and works with the file.
+     */
+    public void useShapeExportedTestWriter(Shape shape, Consumer<GoWriter> writerConsumer) {
+        Symbol symbol = symbolProvider.toSymbol(shape);
+        String filename = symbol.getDefinitionFile();
+
+        StringBuilder b = new StringBuilder(filename);
+        b.insert(filename.lastIndexOf(".go"), "_exported_test");
+        filename = b.toString();
+
+        symbol = symbol.toBuilder()
+                .definitionFile(filename)
+                .namespace(symbol.getNamespace() + "_test", symbol.getNamespaceDelimiter())
+                .build();
 
         useShapeWriter(symbol, writerConsumer);
     }
 
     /**
      * Gets a previously created writer or creates a new one if needed.
      *
-     * @param symbol symbol to create the writer for.
+     * @param symbol         symbol to create the writer for.
      * @param writerConsumer Consumer that accepts and works with the file.
      */
     private void useShapeWriter(Symbol symbol, Consumer<GoWriter> writerConsumer) {
@@ -121,7 +144,7 @@ private void useShapeWriter(Symbol symbol, Consumer<GoWriter> writerConsumer) {
      * Gets a previously created writer or creates a new one if needed
      * and adds a new line if the writer already exists.
      *
-     * @param filename Name of the file to create.
+     * @param filename       Name of the file to create.
      * @param writerConsumer Consumer that accepts and works with the file.
      */
     void useFileWriter(String filename, String namespace, Consumer<GoWriter> writerConsumer) {

codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/UnionGenerator.java

@@ -16,6 +16,8 @@
 package software.amazon.smithy.go.codegen;
 
 import java.util.Collection;
+import java.util.Set;
+import java.util.TreeSet;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.model.Model;
@@ -27,35 +29,44 @@
 /**
  * Renders unions and type aliases for all their members.
  */
-public class UnionGenerator implements Runnable {
+public class UnionGenerator {
     public static final String UNKNOWN_MEMBER_NAME = "UnknownUnionMember";
 
     private final Model model;
     private final SymbolProvider symbolProvider;
-    private final GoWriter writer;
     private final UnionShape shape;
 
-    UnionGenerator(Model model, SymbolProvider symbolProvider, GoWriter writer, UnionShape shape) {
+    UnionGenerator(Model model, SymbolProvider symbolProvider, UnionShape shape) {
         this.model = model;
         this.symbolProvider = symbolProvider;
-        this.writer = writer;
         this.shape = shape;
     }
 
-    @Override
-    public void run() {
+    /**
+     * Generates the Go type definitions for the UnionShape.
+     *
+     * @param writer the writer
+     */
+    public void generateUnion(GoWriter writer) {
         Symbol symbol = symbolProvider.toSymbol(shape);
-        writer.writeShapeDocs(shape);
+        Collection<MemberShape> memberShapes = shape.getAllMembers().values();
 
         // Creates the parent interface for the union, which only defines a
         // non-exported method whose purpose is only to enable satisfying the
         // interface.
+        if (writer.writeShapeDocs(shape)) {
+            writer.writeDocs("");
+        }
+        writer.writeDocs("The following types satisfy this interface:");
+        memberShapes.stream().map(symbolProvider::toMemberName).forEach(name -> {
+            writer.write("//  " + name);
+        });
         writer.openBlock("type $L interface {", "}", symbol.getName(), () -> {
             writer.write("is$L()", symbol.getName());
         }).write("");
 
         // Create structs for each member that satisfy the interface.
-        for (MemberShape member : shape.getAllMembers().values()) {
+        for (MemberShape member : memberShapes) {
             Symbol memberSymbol = symbolProvider.toSymbol(member);
             String exportedMemberName = symbolProvider.toMemberName(member);
             Shape target = model.expectShape(member.getTarget());
@@ -78,6 +89,42 @@ public void run() {
         }
     }
 
+    /**
+     * Generates union usage examples for documentation.
+     *
+     * @param writer the writer
+     */
+    public void generateUnionExamples(GoWriter writer) {
+        Symbol symbol = symbolProvider.toSymbol(shape);
+        Set<MemberShape> members = new TreeSet<>(shape.getAllMembers().values());
+
+        writer.openBlock("func Example$L_outputUsage() {", "}", symbol.getName(), () -> {
+            writer.write("var union $P", symbol);
+
+            writer.writeDocs("type switches can be used to check the union value");
+            writer.openBlock("switch v := union.(type) {", "}", () -> {
+                for (MemberShape member : members) {
+                    Symbol targetSymbol = symbolProvider.toSymbol(model.expectShape(member.getTarget()));
+                    Symbol memberSymbol = SymbolUtils.createValueSymbolBuilder(symbolProvider.toMemberName(member),
+                            symbol.getNamespace()).build();
+
+                    writer.openBlock("case *$T:", "", memberSymbol, () -> {
+                        writer.write("_ = v.Value // Value is $L", targetSymbol.getName());
+                    });
+                }
+                writer.addUseImports(SmithyGoDependency.FMT);
+                Symbol unknownUnionMember = SymbolUtils.createPointableSymbolBuilder("UnknownUnionMember",
+                        symbol.getNamespace()).build();
+                writer.openBlock("case $P:", "", unknownUnionMember, () -> {
+                    writer.write("fmt.Println(\"unknown tag:\", v.Tag)");
+                });
+                writer.openBlock("default:", "", () -> {
+                    writer.write("fmt.Println(\"union is nil or unknown type\")");
+                });
+            });
+        });
+    }
+
     /**
      * Generates a struct for unknown union values that applies to every union in the given set.
      *