Address comment

smithy-lang · Sep 20, 2022 · d7e8572 · d7e8572
1 parent 53bd2ed
commit d7e8572
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/util/Smithy.kt b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/util/Smithy.kt
@@ -82,14 +82,14 @@ fun ServiceShape.hasEventStreamOperations(model: Model): Boolean = operations.an
     model.expectShape(id, OperationShape::class.java).isEventStream(model)
 }
 
-fun MemberShape.redactIfNecessary(model: Model, safeToPrint: String) =
-    model.expectShape(this.target).redactIfNecessary(safeToPrint)
-
-fun Shape.redactIfNecessary(safeToPrint: String) =
-    if (this.hasTrait<SensitiveTrait>()) {
-        "*** Sensitive Data Redacted ***".dq()
-    } else {
-        safeToPrint
+fun Shape.redactIfNecessary(model: Model, safeToPrint: String): String =
+    when (this) {
+        is MemberShape -> model.expectShape(this.target).redactIfNecessary(model, safeToPrint)
+        else -> if (this.hasTrait<SensitiveTrait>()) {
+            "*** Sensitive Data Redacted ***".dq()
+        } else {
+            safeToPrint
+        }
     }
 
 /*