Add import in container shapes

smithy-lang · Nov 4, 2024 · 8bb7957 · 8bb7957
1 parent 00a961a
commit 8bb7957
Showing 1 changed file with 48 additions and 2 deletions.
diff --git a/...codegen/src/main/java/software/amazon/polymorph/smithygo/codegen/ValidationGenerator.java b/...codegen/src/main/java/software/amazon/polymorph/smithygo/codegen/ValidationGenerator.java
@@ -12,6 +12,7 @@
 import software.amazon.polymorph.smithygo.utils.GoCodegenUtils;
 import software.amazon.polymorph.traits.DafnyUtf8BytesTrait;
 import software.amazon.polymorph.traits.ReferenceTrait;
+import software.amazon.smithy.aws.traits.ServiceTrait;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.model.Model;
@@ -492,13 +493,22 @@ private void renderListShape(
           isExternalShape
         );
         if (isExternalShape) {
-          writer.addImportFromModule(
+          if (SmithyNameResolver.isShapeFromAWSSDK(currentShape)) {
+            writer.addImportFromModule(
+              SmithyNameResolver.getGoModuleNameForSdkNamespace(
+                currentShape.getId().getNamespace()
+              ),
+              "types",
+              SmithyNameResolver.smithyTypesNamespace(currentShape)
+            );
+          } else {
+            writer.addImportFromModule(
               SmithyNameResolver.getGoModuleNameForSmithyNamespace(
                 currentShape.getId().getNamespace()
               ),
-              SmithyNameResolver.smithyTypesNamespace(currentShape),
               SmithyNameResolver.smithyTypesNamespace(currentShape)
             );
+          }
         }
         validationFuncInputTypeMap.put(memberShape, inputType);
         dataSourceForList = "Value";
@@ -575,6 +585,24 @@ private void renderMapShape(
           currentShape,
           isExternalShape
         );
+        if (isExternalShape) {
+          if (SmithyNameResolver.isShapeFromAWSSDK(currentShape)) {
+            writer.addImportFromModule(
+              SmithyNameResolver.getGoModuleNameForSdkNamespace(
+                currentShape.getId().getNamespace()
+              ),
+              "types",
+              SmithyNameResolver.smithyTypesNamespace(currentShape)
+            );
+          } else {
+            writer.addImportFromModule(
+              SmithyNameResolver.getGoModuleNameForSmithyNamespace(
+                currentShape.getId().getNamespace()
+              ),
+              SmithyNameResolver.smithyTypesNamespace(currentShape)
+            );
+          }
+        }
         validationFuncInputTypeMap.put(memberShape, inputType);
         dataSourceForMap = "Value";
       }
@@ -626,6 +654,24 @@ private void renderUnionShape(
         currentShape,
         isExternalShape
       );
+      if (isExternalShape) {
+        if (SmithyNameResolver.isShapeFromAWSSDK(currentShape)) {
+          writer.addImportFromModule(
+            SmithyNameResolver.getGoModuleNameForSdkNamespace(
+              currentShape.getId().getNamespace()
+            ),
+            "types",
+            SmithyNameResolver.smithyTypesNamespace(currentShape)
+          );
+        } else {
+          writer.addImportFromModule(
+            SmithyNameResolver.getGoModuleNameForSmithyNamespace(
+              currentShape.getId().getNamespace()
+            ),
+            SmithyNameResolver.smithyTypesNamespace(currentShape)
+          );
+        }
+      }
       validationFuncInputTypeMap.put(memberShape, inputType);
       dataSourceForUnion = "Value";
     }