Handle lists, maps, and unions with flattening

gems/smithy/lib/smithy/model/flattener.rb

@@ -15,7 +15,7 @@ def shape(id)
         shape['mixins'].reverse_each do |mixin|
           mixin_shape = shape(mixin['target'])
           shape = deep_merge(mixin_shape, shape, exclude_traits(mixin_shape))
-          apply_member_traits(id, shape)
+          apply_traits(id, shape)
           shape.delete('mixins')
         end
 
@@ -53,16 +53,62 @@ def deep_merge!(hash1, hash2, exclude_traits, context)
         end
       end
 
-      def apply_member_traits(id, shape)
-        shape.fetch('members', []).each do |member_name, member_shape|
+      def apply_traits(id, shape)
+        case shape['type']
+        when 'structure'
+          structure(id, shape)
+        when 'union'
+          union(id, shape)
+        when 'list'
+          list(id, shape)
+        when 'map'
+          map_key(id, shape)
+          map_value(id, shape)
+        end
+      end
+
+      def structure(id, shape)
+        shape['members'].each do |member_name, member_shape|
           member_id = "#{id}$#{member_name}"
-          next unless @model['shapes'][member_id] && @model['shapes'][member_id]['type'] == 'apply'
+          next unless apply_shape_exists?(member_id)
 
           apply_shape = shape(member_id)
           member_keys = shape['members'][member_name].keys
           shape['members'][member_name] = deep_merge(member_shape, apply_shape).slice(*member_keys)
         end
       end
+      alias union structure
+
+      def list(id, shape)
+        member_id = "#{id}$member"
+        return unless apply_shape_exists?(member_id)
+
+        apply_shape = shape(member_id)
+        member_keys = shape['member'].keys
+        shape['member'] = deep_merge(shape['member'], apply_shape).slice(*member_keys)
+      end
+
+      def map_key(id, shape)
+        key_id = "#{id}$key"
+        return unless apply_shape_exists?(key_id)
+
+        key_shape = shape(key_id)
+        key_keys = shape['key'].keys
+        shape['key'] = deep_merge(shape['key'], key_shape).slice(*key_keys)
+      end
+
+      def map_value(id, shape)
+        value_id = "#{id}$value"
+        return unless apply_shape_exists?(value_id)
+
+        value_shape = shape(value_id)
+        value_keys = shape['value'].keys
+        shape['value'] = deep_merge(shape['value'], value_shape).slice(*value_keys)
+      end
+
+      def apply_shape_exists?(id)
+        @model['shapes'][id] && @model['shapes'][id]['type'] == 'apply'
+      end
     end
   end
 end

gems/smithy/spec/fixtures/mixins/apply-list/model.json

@@ -0,0 +1,31 @@
+{
+    "smithy": "2.0",
+    "shapes": {
+        "smithy.ruby.tests#MyList": {
+            "type": "list",
+            "mixins": [
+                {
+                    "target": "smithy.ruby.tests#MyMixin"
+                }
+            ]
+        },
+        "smithy.ruby.tests#MyList$member": {
+            "type": "apply",
+            "traits": {
+                "smithy.api#documentation": "Specific docs"
+            }
+        },
+        "smithy.ruby.tests#MyMixin": {
+            "type": "list",
+            "member": {
+                "target": "smithy.api#String",
+                "traits": {
+                    "smithy.api#documentation": "Generic docs"
+                }
+            },
+            "traits": {
+                "smithy.api#mixin": {}
+            }
+        }
+    }
+}

gems/smithy/spec/fixtures/mixins/apply-list/model.smithy

@@ -0,0 +1,12 @@
+$version: "2"
+
+namespace smithy.ruby.tests
+
+@mixin
+list MyMixin {
+    /// Generic docs
+    member: String
+}
+
+list MyList with [MyMixin] {}
+apply MyList$member @documentation("Specific docs")

gems/smithy/spec/fixtures/mixins/apply-map/model.json

@@ -0,0 +1,43 @@
+{
+    "smithy": "2.0",
+    "shapes": {
+        "smithy.ruby.tests#MyMap": {
+            "type": "map",
+            "mixins": [
+                {
+                    "target": "smithy.ruby.tests#MyMixin"
+                }
+            ]
+        },
+        "smithy.ruby.tests#MyMap$key": {
+            "type": "apply",
+            "traits": {
+                "smithy.api#documentation": "Specific docs"
+            }
+        },
+        "smithy.ruby.tests#MyMap$value": {
+            "type": "apply",
+            "traits": {
+                "smithy.api#documentation": "Specific docs"
+            }
+        },
+        "smithy.ruby.tests#MyMixin": {
+            "type": "map",
+            "key": {
+                "target": "smithy.api#String",
+                "traits": {
+                    "smithy.api#documentation": "Generic docs"
+                }
+            },
+            "value": {
+                "target": "smithy.api#String",
+                "traits": {
+                    "smithy.api#documentation": "Generic docs"
+                }
+            },
+            "traits": {
+                "smithy.api#mixin": {}
+            }
+        }
+    }
+}

gems/smithy/spec/fixtures/mixins/apply-map/model.smithy

@@ -0,0 +1,15 @@
+$version: "2"
+
+namespace smithy.ruby.tests
+
+@mixin
+map MyMixin {
+    /// Generic docs
+    key: String
+    /// Generic docs
+    value: String
+}
+
+map MyMap with [MyMixin] {}
+apply MyMap$key @documentation("Specific docs")
+apply MyMap$value @documentation("Specific docs")

gems/smithy/spec/fixtures/mixins/apply-union/model.json

@@ -0,0 +1,34 @@
+{
+    "smithy": "2.0",
+    "shapes": {
+        "smithy.ruby.tests#MyMixin": {
+            "type": "union",
+            "members": {
+                "mixinMember": {
+                    "target": "smithy.api#String",
+                    "traits": {
+                        "smithy.api#documentation": "Generic docs"
+                    }
+                }
+            },
+            "traits": {
+                "smithy.api#mixin": {}
+            }
+        },
+        "smithy.ruby.tests#MyUnion": {
+            "type": "union",
+            "mixins": [
+                {
+                    "target": "smithy.ruby.tests#MyMixin"
+                }
+            ],
+            "members": {}
+        },
+        "smithy.ruby.tests#MyUnion$mixinMember": {
+            "type": "apply",
+            "traits": {
+                "smithy.api#documentation": "Specific docs"
+            }
+        }
+    }
+}

gems/smithy/spec/fixtures/mixins/apply-union/model.smithy

@@ -0,0 +1,12 @@
+$version: "2"
+
+namespace smithy.ruby.tests
+
+@mixin
+union MyMixin {
+    /// Generic docs
+    mixinMember: String
+}
+
+union MyUnion with [MyMixin] {}
+apply MyUnion$mixinMember @documentation("Specific docs")

gems/smithy/spec/smithy/model_spec.rb

@@ -145,8 +145,8 @@ module Smithy
           end
         end
 
-        context 'apply' do
-          let(:fixture_name) { 'apply' }
+        context 'apply to structure' do
+          let(:fixture_name) { 'apply-structure' }
 
           it 'resolves mixins' do
             shape = Model.shape(fixture, 'smithy.ruby.tests#MyStruct')
@@ -156,6 +156,41 @@ module Smithy
           end
         end
 
+        context 'apply to union' do
+          let(:fixture_name) { 'apply-union' }
+
+          it 'resolves mixins' do
+            shape = Model.shape(fixture, 'smithy.ruby.tests#MyUnion')
+            expect(shape).to_not have_key('mixins')
+            expect(shape['members']['mixinMember']['traits']['smithy.api#documentation'])
+              .to eq(fixture['shapes']['smithy.ruby.tests#MyUnion$mixinMember']['traits']['smithy.api#documentation'])
+          end
+        end
+
+        context 'apply to list' do
+          let(:fixture_name) { 'apply-list' }
+
+          it 'resolves mixins' do
+            shape = Model.shape(fixture, 'smithy.ruby.tests#MyList')
+            expect(shape).to_not have_key('mixins')
+            expect(shape['member']['traits']['smithy.api#documentation'])
+              .to eq(fixture['shapes']['smithy.ruby.tests#MyList$member']['traits']['smithy.api#documentation'])
+          end
+        end
+
+        context 'apply to map' do
+          let(:fixture_name) { 'apply-map' }
+
+          it 'resolves mixins' do
+            shape = Model.shape(fixture, 'smithy.ruby.tests#MyMap')
+            expect(shape).to_not have_key('mixins')
+            expect(shape['key']['traits']['smithy.api#documentation'])
+              .to eq(fixture['shapes']['smithy.ruby.tests#MyMap$key']['traits']['smithy.api#documentation'])
+            expect(shape['value']['traits']['smithy.api#documentation'])
+              .to eq(fixture['shapes']['smithy.ruby.tests#MyMap$value']['traits']['smithy.api#documentation'])
+          end
+        end
+
         context 'redefine' do
           let(:fixture_name) { 'redefine' }