From cae26c510907b13d2afe44bc4ad23aa3cc70d1ac Mon Sep 17 00:00:00 2001
From: 0marperez <operez899999@gmail.com>
Date: Fri, 10 Nov 2023 04:17:08 -0500
Subject: [PATCH] function to find if operation payload is union

---
 .../amazon/smithy/kotlin/codegen/model/ShapeExt.kt       | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/model/ShapeExt.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/model/ShapeExt.kt
index 3ff18b28b..a65402bc2 100644
--- a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/model/ShapeExt.kt
+++ b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/model/ShapeExt.kt
@@ -192,6 +192,15 @@ val Shape.isSparse: Boolean
 val Shape.isStreaming: Boolean
     get() = hasTrait<StreamingTrait>()
 
+/**
+ * Returns boolean indicating if operation input is union shaped
+ */
+fun OperationShape.payloadIsUnionShape(model: Model): Boolean {
+    val requestShape = model.expectShape<StructureShape>(input.get())
+    val payload = requestShape.findMemberWithTrait<HttpPayloadTrait>(model)?.targetOrSelf(model)
+    return payload is UnionShape
+}
+
 /**
  * Test if a member targets an event stream
  */