add support for transforms in codegen

disneystreaming · Jan 14, 2022 · 05ab381 · 05ab381
1 parent 1fbbe67
commit 05ab381
Show file tree

Hide file tree

Showing 11 changed files with 66 additions and 17 deletions.
diff --git a/build.sbt b/build.sbt
@@ -303,6 +303,7 @@ lazy val codegen = projectMatrix
     libraryDependencies ++= Seq(
       Dependencies.Cats.core.value,
       Dependencies.Smithy.model,
+      Dependencies.Smithy.build,
       Dependencies.Smithy.awsTraits,
       Dependencies.Smithy.waiters,
       "com.lihaoyi" %% "os-lib" % "0.8.0",
@@ -591,6 +592,7 @@ lazy val Dependencies = new {
   val Smithy = new {
     val smithyVersion = "1.16.2"
     val model = "software.amazon.smithy" % "smithy-model" % smithyVersion
+    val build = "software.amazon.smithy" % "smithy-build" % smithyVersion
     val awsTraits =
       "software.amazon.smithy" % "smithy-aws-traits" % smithyVersion
     val openapi = "software.amazon.smithy" % "smithy-openapi" % smithyVersion

diff --git a/modules/codegen-cli/src/smithy4s/codegen/cli/CodegenCommand.scala b/modules/codegen-cli/src/smithy4s/codegen/cli/CodegenCommand.scala
@@ -87,11 +87,12 @@ object CodegenCommand {
       allowedNSOpt,
       repositoriesOpt,
       dependenciesOpt,
+      transformersOpt,
       specsArgs
     )
       .mapN {
         // format: off
-        case (output, openApiOutput, skipScala, skipOpenapi, allowedNS, repositories, dependencies, specsArgs) =>
+        case (output, openApiOutput, skipScala, skipOpenapi, allowedNS, repositories, dependencies, transformers, specsArgs) =>
         // format: on
           CodegenArgs(
             specsArgs,
@@ -101,7 +102,8 @@ object CodegenCommand {
             skipOpenapi,
             allowedNS,
             repositories.getOrElse(List.empty),
-            dependencies.getOrElse(List.empty)
+            dependencies.getOrElse(List.empty),
+            transformers.getOrElse(List.empty)
           )
       }
 

diff --git a/modules/codegen-cli/src/smithy4s/codegen/cli/DumpModel.scala b/modules/codegen-cli/src/smithy4s/codegen/cli/DumpModel.scala
@@ -25,7 +25,8 @@ object DumpModel {
     val (_, model) = ModelLoader.load(
       args.specs.map(_.toIO).toSet,
       args.dependencies,
-      args.repositories
+      args.repositories,
+      args.transformers
     )
 
     Node.prettyPrintJson(ModelSerializer.builder().build.serialize(model))

diff --git a/modules/codegen-cli/src/smithy4s/codegen/cli/DumpModelCommand.scala b/modules/codegen-cli/src/smithy4s/codegen/cli/DumpModelCommand.scala
@@ -27,7 +27,8 @@ object DumpModelCommand {
   val options = (
     specsArgs,
     repositoriesOpt.map(_.getOrElse(Nil)),
-    dependenciesOpt.map(_.getOrElse(Nil))
+    dependenciesOpt.map(_.getOrElse(Nil)),
+    transformersOpt.map(_.getOrElse(Nil))
   ).mapN(DumpModelArgs.apply)
 
   val command: Command[DumpModel] =

diff --git a/modules/codegen-cli/src/smithy4s/codegen/cli/Options.scala b/modules/codegen-cli/src/smithy4s/codegen/cli/Options.scala
@@ -72,4 +72,13 @@ object Options {
       .map(_.split(',').toList)
       .orNone
 
+  val transformersOpt: Opts[Option[List[String]]] =
+    Opts
+      .option[String](
+        "transformers",
+        "Comma-delimited list of transformer names to apply to smithy files"
+      )
+      .map(_.split(',').toList)
+      .orNone
+
 }
diff --git a/modules/codegen-cli/src/smithy4s/codegen/cli/Smithy4sCommand.scala b/modules/codegen-cli/src/smithy4s/codegen/cli/Smithy4sCommand.scala
@@ -26,7 +26,8 @@ object Smithy4sCommand {
   final case class DumpModelArgs(
       specs: List[os.Path],
       repositories: List[String],
-      dependencies: List[String]
+      dependencies: List[String],
+      transformers: List[String]
   )
 
 }
diff --git a/modules/codegen-cli/test/src/smithy4s/codegen/cli/CommandParsingSpec.scala b/modules/codegen-cli/test/src/smithy4s/codegen/cli/CommandParsingSpec.scala
@@ -33,7 +33,8 @@ object CommandParsingSpec extends FunSuite {
               skipOpenapi = false,
               allowedNS = None,
               repositories = Nil,
-              dependencies = Nil
+              dependencies = Nil,
+              transformers = Nil
             )
           )
         )
@@ -56,7 +57,9 @@ object CommandParsingSpec extends FunSuite {
         "--repositories",
         "repo1,repo2",
         "--dependencies",
-        "dep1,dep2"
+        "dep1,dep2",
+        "--transformers",
+        "t1,t2"
       )
     )
 
@@ -75,7 +78,8 @@ object CommandParsingSpec extends FunSuite {
               skipOpenapi = true,
               allowedNS = Some(Set("name1", "name2")),
               repositories = List("repo1", "repo2"),
-              dependencies = List("dep1", "dep2")
+              dependencies = List("dep1", "dep2"),
+              transformers = List("t1", "t2")
             )
           )
         )
@@ -90,7 +94,8 @@ object CommandParsingSpec extends FunSuite {
             Smithy4sCommand.DumpModelArgs(
               specs = Nil,
               repositories = Nil,
-              dependencies = Nil
+              dependencies = Nil,
+              transformers = Nil
             )
           )
         )
@@ -105,7 +110,9 @@ object CommandParsingSpec extends FunSuite {
         "--repositories",
         "repo1,repo2",
         "--dependencies",
-        "dep1,dep2"
+        "dep1,dep2",
+        "--transformers",
+        "t1,t2"
       )
     )
     assert(
@@ -118,7 +125,8 @@ object CommandParsingSpec extends FunSuite {
                 os.pwd / "sampleSpecs" / "example.smithy"
               ),
               repositories = List("repo1", "repo2"),
-              dependencies = List("dep1", "dep2")
+              dependencies = List("dep1", "dep2"),
+              transformers = List("t1", "t2")
             )
           )
         )

diff --git a/modules/codegen-plugin/src/smithy4s/codegen/CodegenPlugin.scala b/modules/codegen-plugin/src/smithy4s/codegen/CodegenPlugin.scala
@@ -51,6 +51,11 @@ object Smithy4sCodegenPlugin extends AutoPlugin {
       settingKey[List[String]](
         "List of dependencies containing smithy files to include in codegen task"
       )
+
+    val smithy4sModelTransformers =
+      settingKey[List[String]](
+        "List of transformer names that should be applied to the model prior to codegen"
+      )
   }
 
   import autoImport._
@@ -72,7 +77,8 @@ object Smithy4sCodegenPlugin extends AutoPlugin {
       Compile / resourceGenerators += (Compile / smithy4sCodegen).map(
         _.filter(_.ext != "scala")
       ),
-      cleanFiles += (Compile / smithy4sOutputDir).value
+      cleanFiles += (Compile / smithy4sOutputDir).value,
+      Compile / smithy4sModelTransformers := List.empty
     )
 
   private def untupled[A, B, C](f: ((A, B)) => C): (A, B) => C = (a, b) =>
@@ -96,6 +102,7 @@ object Smithy4sCodegenPlugin extends AutoPlugin {
       (conf / resolvers).value.toList.collect { case m: MavenRepository =>
         m.root
       }
+    val transforms = (conf / smithy4sModelTransformers).value
     val s = streams.value
 
     val cached =
@@ -117,7 +124,8 @@ object Smithy4sCodegenPlugin extends AutoPlugin {
                   skipOpenapi = false,
                   allowedNS = allowedNamespaces,
                   repositories = res,
-                  dependencies = dependencies
+                  dependencies = dependencies,
+                  transforms
                 )
                 val resPaths = smithy4s.codegen.Codegen
                   .processSpecs(codegenArgs)

diff --git a/modules/codegen/src/smithy4s/codegen/Codegen.scala b/modules/codegen/src/smithy4s/codegen/Codegen.scala
@@ -30,7 +30,8 @@ object Codegen { self =>
     val (classloader, model) = ModelLoader.load(
       args.specs.map(_.toIO).toSet,
       args.dependencies,
-      args.repositories
+      args.repositories,
+      args.transformers
     )
 
     val scalaFiles = if (!args.skipScala) {

diff --git a/modules/codegen/src/smithy4s/codegen/CodegenArgs.scala b/modules/codegen/src/smithy4s/codegen/CodegenArgs.scala
@@ -24,5 +24,6 @@ case class CodegenArgs(
     skipOpenapi: Boolean,
     allowedNS: Option[Set[String]],
     repositories: List[String],
-    dependencies: List[String]
+    dependencies: List[String],
+    transformers: List[String]
 )
diff --git a/modules/codegen/src/smithy4s/codegen/ModelLoader.scala b/modules/codegen/src/smithy4s/codegen/ModelLoader.scala
@@ -24,12 +24,15 @@ import software.amazon.smithy.model.Model
 import java.io.File
 import java.net.URL
 import java.net.URLClassLoader
+import software.amazon.smithy.build.ProjectionTransformer
+import software.amazon.smithy.build.TransformContext
 
 object ModelLoader {
   def load(
       specs: Set[File],
       dependencies: List[String],
-      repositories: List[String]
+      repositories: List[String],
+      transformers: List[String]
   ): (ClassLoader, Model) = {
     val maybeDeps = resolveDependencies(dependencies, repositories)
     val currentClassLoader = this.getClass().getClassLoader()
@@ -76,7 +79,19 @@ object ModelLoader {
       .assemble()
       .unwrap()
 
-    (validatorClassLoader, model)
+    val serviceFactory =
+      ProjectionTransformer.createServiceFactory(validatorClassLoader)
+
+    val trans = transformers.flatMap { t =>
+      val result = serviceFactory(t)
+      if (result.isPresent()) Some(result.get) else None
+    }
+
+    val transformedModel = trans.foldLeft(model)((m, t) =>
+      t.transform(TransformContext.builder().model(m).build())
+    )
+
+    (validatorClassLoader, transformedModel)
   }
 
   private def resolveDependencies(