The world if Smithy4s had the classloader fix

project/plugins.sbt

@@ -1,3 +1,3 @@
-addSbtPlugin("com.disneystreaming.smithy4s" % "smithy4s-sbt-codegen" % "0.18.35")
+addSbtPlugin("com.disneystreaming.smithy4s" % "smithy4s-sbt-codegen" % "0.18.36")
 addSbtPlugin("org.scala-native" % "sbt-scala-native" % "0.4.17")
 addSbtPlugin("org.typelevel" % "sbt-typelevel" % "0.7.7")

transformation/src/main/scala/OpenEnums.scala

@@ -15,50 +15,29 @@
  */
 
 import alloy.OpenEnumTrait
-import bsp.traits.EnumKindTrait
 import bsp.traits.EnumKindTrait.EnumKind.CLOSED
 import bsp.traits.EnumKindTrait.EnumKind.OPEN
+import common.TransformationUtils.*
 import software.amazon.smithy.build.ProjectionTransformer
 import software.amazon.smithy.build.TransformContext
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.AbstractShapeBuilder
 import software.amazon.smithy.model.shapes.Shape
-import software.amazon.smithy.model.transform.ModelTransformer
-
-import scala.collection.JavaConverters.*
 
 class OpenEnums extends ProjectionTransformer {
   def getName(): String = "open-enums"
 
-  def transform(context: TransformContext): Model = ModelTransformer
-    .create()
-    .mapShapes(
-      context.getModel(),
-      s =>
-        s match {
-          case s if s.hasTrait(bsp.traits.EnumKindTrait.ID) =>
-            val builder = Shape.shapeToBuilder(s): AbstractShapeBuilder[_, _]
-            builder.removeTrait(bsp.traits.EnumKindTrait.ID)
-
-            val dynTrait = s
-              .getAllTraits()
-              .asScala
-              .apply(bsp.traits.EnumKindTrait.ID)
-
-            new EnumKindTrait.Provider()
-              .createTrait(
-                dynTrait.toShapeId(),
-                dynTrait.toNode(),
-              )
-              .getEnumKind() match {
-              case OPEN   => builder.addTrait(new OpenEnumTrait())
-              case CLOSED => () // do nothing, we remove the trait anyway
-            }
+  def transform(context: TransformContext): Model = context.getModel().mapSomeShapes {
+    case s if s.hasTrait(bsp.traits.EnumKindTrait.ID) =>
+      val builder = Shape.shapeToBuilder(s): AbstractShapeBuilder[_, _]
+      builder.removeTrait(bsp.traits.EnumKindTrait.ID)
 
-            builder.build()
+      s.expectTrait(classOf[bsp.traits.EnumKindTrait]).getEnumKind() match {
+        case OPEN   => builder.addTrait(new OpenEnumTrait())
+        case CLOSED => () // do nothing, we remove the trait anyway
+      }
 
-          case s => s
-        },
-    )
+      builder.build()
+  }
 
 }

transformation/src/main/scala/SetShapes.scala

@@ -14,28 +14,18 @@
  * limitations under the License.
  */
 
+import bsp.traits.SetTrait
+import common.TransformationUtils.*
 import software.amazon.smithy.build.ProjectionTransformer
 import software.amazon.smithy.build.TransformContext
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.traits.UniqueItemsTrait
-import software.amazon.smithy.model.transform.ModelTransformer
 
 class SetShapes extends ProjectionTransformer {
   def getName(): String = "set-shapes"
 
-  def transform(context: TransformContext): Model = ModelTransformer
-    .create()
-    .mapShapes(
-      context.getModel(),
-      s =>
-        s match {
-          case s if s.hasTrait(bsp.traits.SetTrait.ID) =>
-            val builder = s.asListShape().get().toBuilder()
-            builder.removeTrait(bsp.traits.SetTrait.ID)
-            builder.addTrait(new UniqueItemsTrait())
-            builder.build()
-          case s => s
-        },
-    )
+  def transform(context: TransformContext): Model = context.getModel.mapSomeTraits {
+    case (_, _: SetTrait) => new UniqueItemsTrait()
+  }
 
 }

transformation/src/main/scala/TransformBuildTargetData.scala

@@ -347,16 +347,7 @@ class TransformBuildTargetData extends ProjectionTransformer {
     else
       None
 
-  private def expectDataKind(s: Shape): DataKindTrait = {
-    val trt = Option(
-      s.getAllTraits().get(DataKindTrait.ID)
-    ).getOrElse(sys.error(s"Expected $s to have a DataKind trait, but it doesn't"))
-
-    new DataKindTrait.Provider().createTrait(
-      trt.toShapeId(),
-      trt.toNode(),
-    )
-  }
+  private def expectDataKind(s: Shape): DataKindTrait = s.expectTrait(classOf[DataKindTrait])
 
   // for debugging modified smithy
   private def dump(m: Model): Unit = {

transformation/src/main/scala/TransformJsonRpcTraits.scala

@@ -14,52 +14,25 @@
  * limitations under the License.
  */
 
+import bsp.traits.JsonNotificationTrait
+import bsp.traits.JsonRPCTrait
+import bsp.traits.JsonRequestTrait
+import common.TransformationUtils.*
 import software.amazon.smithy.build.ProjectionTransformer
 import software.amazon.smithy.build.TransformContext
 import software.amazon.smithy.model.Model
-import software.amazon.smithy.model.transform.ModelTransformer
-import bsp.traits.JsonNotificationTrait
-import bsp.traits.JsonRequestTrait
-import bsp.traits.JsonRPCTrait
 
 class TransformJsonRpcTraits extends ProjectionTransformer {
   def getName(): String = "transform-jsonrpclib-traits"
 
-  def transform(context: TransformContext): Model = ModelTransformer
-    .create()
-    .mapShapes(
-      context.getModel(),
-      s =>
-        s match {
-          case s if s.hasTrait(JsonRPCTrait.ID) =>
-            val builder = s.asServiceShape.get.toBuilder()
-            builder.removeTrait(JsonRPCTrait.ID)
-            builder.addTrait(jsonrpclib.JsonRPCTrait.builder().build())
-            builder.build()
-
-          case s if s.hasTrait(JsonNotificationTrait.ID) =>
-            val builder = s.asOperationShape.get.toBuilder()
-            builder.removeTrait(JsonNotificationTrait.ID)
-            builder.addTrait(
-              new jsonrpclib.JsonNotificationTrait.Provider().createTrait(
-                jsonrpclib.JsonNotificationTrait.ID,
-                s.getAllTraits().get(JsonNotificationTrait.ID).toNode(),
-              )
-            )
-            builder.build()
-
-          case s if s.hasTrait(JsonRequestTrait.ID) =>
-            val builder = s.asOperationShape.get.toBuilder()
-            builder.removeTrait(JsonRequestTrait.ID)
-            builder.addTrait(
-              new jsonrpclib.JsonRequestTrait.Provider().createTrait(
-                jsonrpclib.JsonRequestTrait.ID,
-                s.getAllTraits().get(JsonRequestTrait.ID).toNode(),
-              )
-            )
-            builder.build()
-          case s => s
-        },
+  def transform(context: TransformContext): Model = context
+    .getModel()
+    .mapSomeTraits(
+      { case (_, _: JsonRPCTrait) => jsonrpclib.JsonRPCTrait.builder().build() },
+      { case (_, trt: JsonNotificationTrait) =>
+        new jsonrpclib.JsonNotificationTrait(trt.getValue())
+      },
+      { case (_, trt: JsonRequestTrait) => new jsonrpclib.JsonRequestTrait(trt.getValue()) },
     )
 
 }

transformation/src/main/scala/UntaggedUnions.scala

@@ -14,27 +14,16 @@
  * limitations under the License.
  */
 
+import common.TransformationUtils.*
 import software.amazon.smithy.build.ProjectionTransformer
 import software.amazon.smithy.build.TransformContext
 import software.amazon.smithy.model.Model
-import software.amazon.smithy.model.transform.ModelTransformer
 
 class UntaggedUnions extends ProjectionTransformer {
   def getName(): String = "untagged-unions"
 
-  def transform(context: TransformContext): Model = ModelTransformer
-    .create()
-    .mapShapes(
-      context.getModel(),
-      s =>
-        s match {
-          case s if s.hasTrait(bsp.traits.UntaggedUnionTrait.ID) =>
-            val builder = s.asUnionShape().get().toBuilder()
-            builder.removeTrait(bsp.traits.UntaggedUnionTrait.ID)
-            builder.addTrait(new alloy.UntaggedUnionTrait())
-            builder.build()
-          case s => s
-        },
-    )
+  def transform(context: TransformContext): Model = context.getModel().mapSomeTraits {
+    case (_, _: bsp.traits.UntaggedUnionTrait) => new alloy.UntaggedUnionTrait()
+  }
 
 }

transformation/src/main/scala/common/TransformationUtils.scala

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2025 Polyvariant
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package common
+
+import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.transform.ModelTransformer
+import software.amazon.smithy.model.shapes.Shape
+import software.amazon.smithy.model.traits.Trait
+import java.util.function.BiFunction
+import scala.collection.JavaConverters.*
+
+object TransformationUtils {
+
+  implicit class ModelOps(private val model: Model) extends AnyVal {
+
+    def mapSomeShapes(f: PartialFunction[Shape, Shape]): Model = ModelTransformer
+      .create()
+      .mapShapes(model, pfOrIdentity(f).apply(_))
+
+    def mapSomeTraits(funs: PartialFunction[(Shape, Trait), Trait]*): Model = ModelTransformer
+      .create()
+      .mapTraits(
+        model,
+        funs
+          .map(pf =>
+            ((s, trt) => pf.lift((s, trt)).getOrElse(trt)): BiFunction[Shape, Trait, Trait]
+          )
+          .asJava,
+      )
+
+  }
+
+  def pfOrIdentity[A](f: PartialFunction[A, A]): A => A = a => f.lift(a).getOrElse(a)
+}