diff --git a/project/plugins.sbt b/project/plugins.sbt index d910113..ef9b9cf 100644 --- a/project/plugins.sbt +++ b/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") diff --git a/transformation/src/main/scala/OpenEnums.scala b/transformation/src/main/scala/OpenEnums.scala index f0ecb55..1bf1e5e 100644 --- a/transformation/src/main/scala/OpenEnums.scala +++ b/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() + } } diff --git a/transformation/src/main/scala/SetShapes.scala b/transformation/src/main/scala/SetShapes.scala index 016071f..fe3bd24 100644 --- a/transformation/src/main/scala/SetShapes.scala +++ b/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() + } } diff --git a/transformation/src/main/scala/TransformBuildTargetData.scala b/transformation/src/main/scala/TransformBuildTargetData.scala index e793b7f..be38597 100644 --- a/transformation/src/main/scala/TransformBuildTargetData.scala +++ b/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 = { diff --git a/transformation/src/main/scala/TransformJsonRpcTraits.scala b/transformation/src/main/scala/TransformJsonRpcTraits.scala index 6398d95..e1a9526 100644 --- a/transformation/src/main/scala/TransformJsonRpcTraits.scala +++ b/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()) }, ) } diff --git a/transformation/src/main/scala/UntaggedUnions.scala b/transformation/src/main/scala/UntaggedUnions.scala index 0f6adad..8e5a1c6 100644 --- a/transformation/src/main/scala/UntaggedUnions.scala +++ b/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() + } } diff --git a/transformation/src/main/scala/common/TransformationUtils.scala b/transformation/src/main/scala/common/TransformationUtils.scala new file mode 100644 index 0000000..2eafd87 --- /dev/null +++ b/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) +}