Add protocol alias trait protocol

To differentiate between alias traits attached to headers and alias traits specific to a protocol
soto-project · Dec 4, 2023 · b2574a0 · b2574a0
1 parent a075f3d
commit b2574a0
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 16 deletions.
diff --git a/Sources/SotoCodeGeneratorLib/AwsService+shapes.swift b/Sources/SotoCodeGeneratorLib/AwsService+shapes.swift
@@ -333,7 +333,7 @@ extension AwsService {
             // if part of URL
         } else if member.hasTrait(type: HttpLabelTrait.self) {
             let labelName = isPropertyWrapper ? "_\(name.toSwiftLabelCase())" : name.toSwiftLabelCase()
-            let aliasTrait = member.trait(named: serviceProtocolTrait.nameTrait.staticName) as? AliasTrait
+            let aliasTrait = member.trait(named: serviceProtocolTrait.nameTrait.staticName) as? ProtocolAliasTrait
             memberEncoding.append(.init(name: labelName, location: ".uri(\"\(aliasTrait?.alias ?? name)\")"))
             // if response status code
         } else if member.hasTrait(type: HttpResponseCodeTrait.self) {
@@ -343,7 +343,7 @@ extension AwsService {
         } else if member.hasTrait(type: HttpPayloadTrait.self),
                   !(model.shape(for: member.target) is BlobShape) || isOutputShape
         {
-            let aliasTrait = member.traits?.first(where: { $0 is AliasTrait }) as? AliasTrait
+            let aliasTrait = member.trait(named: serviceProtocolTrait.nameTrait.staticName) as? ProtocolAliasTrait
             let payloadName = aliasTrait?.alias ?? name
             let swiftLabelName = name.toSwiftLabelCase()
             if swiftLabelName != payloadName {
@@ -354,7 +354,7 @@ extension AwsService {
 
         if member.hasTrait(type: HostLabelTrait.self) {
             let labelName = isPropertyWrapper ? "_\(name.toSwiftLabelCase())" : name.toSwiftLabelCase()
-            let aliasTrait = member.trait(named: serviceProtocolTrait.nameTrait.staticName) as? AliasTrait
+            let aliasTrait = member.trait(named: serviceProtocolTrait.nameTrait.staticName) as? ProtocolAliasTrait
             memberEncoding.append(.init(name: labelName, location: ".hostname(\"\(aliasTrait?.alias ?? name)\")"))
         }
         return memberEncoding
@@ -371,7 +371,9 @@ extension AwsService {
             return nil
         }
         var rawValue: String = name
-        if let aliasTrait = member.traits?.first(where: { $0 is AliasTrait }) as? AliasTrait {
+        if let aliasTrait = member.trait(named: serviceProtocolTrait.nameTrait.staticName) as? ProtocolAliasTrait {
+            rawValue = aliasTrait.alias
+        } else if let aliasTrait = member.traits?.first(where: { $0 is AliasTrait }) as? AliasTrait {
             rawValue = aliasTrait.alias
         }
         let variable = name.toSwiftVariableCase()

diff --git a/Sources/SotoCodeGeneratorLib/AwsService.swift b/Sources/SotoCodeGeneratorLib/AwsService.swift
@@ -288,14 +288,14 @@ struct AwsService {
 
     func getListEntryName(member: MemberShape, list: ListShape) -> String? {
         guard !member.hasTrait(type: XmlFlattenedTrait.self) else { return nil }
-        guard let memberName = list.member.traits?.first(where: { $0 is AliasTrait }) as? AliasTrait else { return "member" }
+        guard let memberName = list.member.trait(named: serviceProtocolTrait.nameTrait.staticName) as? ProtocolAliasTrait else { return "member" }
         return memberName.alias
     }
 
     func getMapEntryNames(member: MemberShape, map: MapShape) -> (entry: String?, key: String, value: String) {
         let flattened = member.hasTrait(type: XmlFlattenedTrait.self)
-        let keyTrait = map.key.traits?.first(where: { $0 is AliasTrait }) as? AliasTrait
-        let valueTrait = map.value.traits?.first(where: { $0 is AliasTrait }) as? AliasTrait
+        let keyTrait = map.key.trait(named: self.serviceProtocolTrait.nameTrait.staticName) as? ProtocolAliasTrait
+        let valueTrait = map.value.trait(named: self.serviceProtocolTrait.nameTrait.staticName) as? ProtocolAliasTrait
         return (entry: flattened ? nil : "entry", key: keyTrait?.alias ?? "key", value: valueTrait?.alias ?? "value")
     }
 

diff --git a/Sources/SotoCodeGeneratorLib/Smithy+CodeGeneration.swift b/Sources/SotoCodeGeneratorLib/Smithy+CodeGeneration.swift
@@ -23,15 +23,19 @@ protocol AliasTrait: StaticTrait {
     var alias: String { get }
 }
 
+protocol ProtocolAliasTrait: StaticTrait {
+    var alias: String { get }
+}
+
 protocol AwsServiceProtocol: SotoOutput {
-    var nameTrait: AliasTrait.Type { get }
+    var nameTrait: ProtocolAliasTrait.Type { get }
     var requiresCollectionCoders: Bool { get }
 }
 
 // MARK: Alias traits
 
-extension JsonNameTrait: AliasTrait { var alias: String { return value } }
-extension XmlNameTrait: AliasTrait { var alias: String { return value } }
+extension JsonNameTrait: ProtocolAliasTrait { var alias: String { return value } }
+extension XmlNameTrait: ProtocolAliasTrait { var alias: String { return value } }
 // going to have to assume EC2 name isn't used and use Header name
 // extension AwsProtocolsEc2QueryNameTrait: AliasTrait { var alias: String { return value } }
 extension HttpHeaderTrait: AliasTrait { var alias: String { return value } }
@@ -42,37 +46,37 @@ extension HttpQueryTrait: AliasTrait { var alias: String { return value } }
 
 extension AwsProtocolsRestJson1Trait: AwsServiceProtocol {
     var output: String { ".restjson" }
-    var nameTrait: AliasTrait.Type { return JsonNameTrait.self }
+    var nameTrait: ProtocolAliasTrait.Type { return JsonNameTrait.self }
     var requiresCollectionCoders: Bool { return false }
 }
 
 extension AwsProtocolsAwsJson1_0Trait: AwsServiceProtocol {
     var output: String { ".json(version: \"1.0\")" }
-    var nameTrait: AliasTrait.Type { return JsonNameTrait.self }
+    var nameTrait: ProtocolAliasTrait.Type { return JsonNameTrait.self }
     var requiresCollectionCoders: Bool { return false }
 }
 
 extension AwsProtocolsAwsJson1_1Trait: AwsServiceProtocol {
     var output: String { ".json(version: \"1.1\")" }
-    var nameTrait: AliasTrait.Type { return JsonNameTrait.self }
+    var nameTrait: ProtocolAliasTrait.Type { return JsonNameTrait.self }
     var requiresCollectionCoders: Bool { return false }
 }
 
 extension AwsProtocolsAwsQueryTrait: AwsServiceProtocol {
     var output: String { ".query" }
-    var nameTrait: AliasTrait.Type { return XmlNameTrait.self }
+    var nameTrait: ProtocolAliasTrait.Type { return XmlNameTrait.self }
     var requiresCollectionCoders: Bool { return true }
 }
 
 extension AwsProtocolsEc2QueryTrait: AwsServiceProtocol {
     var output: String { ".ec2" }
-    var nameTrait: AliasTrait.Type { return XmlNameTrait.self }
+    var nameTrait: ProtocolAliasTrait.Type { return XmlNameTrait.self }
     var requiresCollectionCoders: Bool { return true }
 }
 
 extension AwsProtocolsRestXmlTrait: AwsServiceProtocol {
     var output: String { ".restxml" }
-    var nameTrait: AliasTrait.Type { return XmlNameTrait.self }
+    var nameTrait: ProtocolAliasTrait.Type { return XmlNameTrait.self }
     var requiresCollectionCoders: Bool { return true }
 }