Add custom key shape for maps

Author: Nicholas Thomson

pkg/sdk/custom_shapes.go

@@ -27,6 +27,7 @@ var (
 const (
 	ShapeNameTemplateList = "%sList"
 	ShapeNameTemplateMap  = "%sMap"
+	ShapeNameTemplateKey  = "%sKey"
 )
 
 type CustomShape struct {
@@ -113,6 +114,16 @@ func (i *customShapeInjector) createShapeRefForMember(shape *awssdkmodel.Shape)
 	}
 }
 
+// createKeyShape creates a Shape that acts as the string key shape for a
+// custom map.
+func (i *customShapeInjector) createKeyShape(shapeName string) *awssdkmodel.Shape {
+	return &awssdkmodel.Shape{
+		API:       i.sdkAPI.API,
+		ShapeName: fmt.Sprintf(ShapeNameTemplateKey, shapeName),
+		Type:      "string",
+	}
+}
+
 // newList loads a shape given its name and creates a custom shape that is a
 // map with strings as keys and that shape as the value.
 func (i *customShapeInjector) newMap(valueShapeName string) (*CustomShape, error) {
@@ -125,11 +136,15 @@ func (i *customShapeInjector) newMap(valueShapeName string) (*CustomShape, error
 	shapeName := fmt.Sprintf(ShapeNameTemplateMap, valueShape.ShapeName)
 	documentation := ""
 
+	keyShape := i.createKeyShape(shapeName)
+	keyShapeRef := i.createShapeRefForMember(keyShape)
+
 	shape := &awssdkmodel.Shape{
 		API:       i.sdkAPI.API,
 		ShapeName: shapeName,
 		// TODO (RedbackThomson): Support documentation for custom shapes
 		Documentation: documentation,
+		KeyRef:        *keyShapeRef,
 		ValueRef:      *valueShapeRef,
 		Type:          "map",
 	}