Don't crash on no input, allow customizing source directory, support external trait providers (#2)

* First couple improvements

* Support external trait providers

* Add trailing newline

* break mima

* Add hack

Author: kubukoz

build.sbt

@@ -1,4 +1,4 @@
-ThisBuild / tlBaseVersion := "0.1"
+ThisBuild / tlBaseVersion := "0.2"
 ThisBuild / organization := "org.polyvariant"
 ThisBuild / organizationName := "Polyvariant"
 ThisBuild / startYear := Some(2025)

sbtPlugin/src/main/scala/org/polyvariant/smithytraitcodegen/SmithyTraitCodegen.scala

@@ -25,6 +25,9 @@ import software.amazon.smithy.model.node.ObjectNode
 import software.amazon.smithy.traitcodegen.TraitCodegenPlugin
 
 import java.io.File
+import software.amazon.smithy.model.transform.ModelTransformer
+import java.util.stream.Collectors
+import scala.collection.JavaConverters.*
 
 object SmithyTraitCodegen {
 
@@ -38,12 +41,13 @@ object SmithyTraitCodegen {
     targetDir: os.Path,
     smithySourcesDir: PathRef,
     dependencies: List[PathRef],
+    externalProviders: List[String],
   )
 
   object Args {
 
     // format: off
-    private type ArgsDeconstructed = String :*: String :*: os.Path :*: PathRef :*: List[PathRef] :*: LNil
+    private type ArgsDeconstructed = String :*: String :*: os.Path :*: PathRef :*: List[PathRef] :*: List[String] :*: LNil
     // format: on
 
     private implicit val pathFormat: JsonFormat[os.Path] = BasicJsonProtocol
@@ -56,6 +60,7 @@ object SmithyTraitCodegen {
           ("targetDir", args.targetDir) :*:
           ("smithySourcesDir", args.smithySourcesDir) :*:
           ("dependencies", args.dependencies) :*:
+          ("externalProviders", args.externalProviders) :*:
           LNil
       },
       {
@@ -64,13 +69,15 @@ object SmithyTraitCodegen {
             (_, targetDir) :*:
             (_, smithySourcesDir) :*:
             (_, dependencies) :*:
+            (_, externalProviders) :*:
             LNil =>
           Args(
             javaPackage = javaPackage,
             smithyNamespace = smithyNamespace,
             targetDir = targetDir,
             smithySourcesDir = smithySourcesDir,
             dependencies = dependencies,
+            externalProviders = externalProviders,
           )
       },
     )
@@ -104,6 +111,23 @@ object SmithyTraitCodegen {
 
   }
 
+  // Hack / workaround for https://github.com/smithy-lang/smithy/pull/2671
+  private def namespaceHackRequired(ns: String) =
+    ns.startsWith("smithy") && ns != "smithy" && !ns.startsWith("smithy.")
+
+  private def renameNamespaceForHack(ns: String) =
+    if (namespaceHackRequired(ns))
+      "hack" + ns
+    else
+      ns
+
+  private def replaceNamespaceRefsInFile(fileText: String, ns: String) =
+    if (namespaceHackRequired(ns)) {
+      fileText.replaceAll(s"hack$ns", ns)
+    } else {
+      fileText
+    }
+
   def generate(args: Args): Output = {
     val outputDir = args.targetDir / "smithy-trait-generator-output"
     val genDir = outputDir / "java"
@@ -113,13 +137,38 @@ object SmithyTraitCodegen {
 
     val manifest = FileManifest.create(genDir.toNIO)
 
-    val model = args
-      .dependencies
-      .foldLeft(Model.assembler().addImport(args.smithySourcesDir.path.toNIO)) { case (acc, dep) =>
-        acc.addImport(dep.path.toNIO)
+    val model =
+      args
+        .dependencies
+        .foldLeft(Model.assembler().addImport(args.smithySourcesDir.path.toNIO)) {
+          case (acc, dep) => acc.addImport(dep.path.toNIO)
+        }
+        .assemble()
+        .unwrap() match {
+        case model =>
+          if (namespaceHackRequired(args.smithyNamespace)) {
+            println("Applying namespace workaround - `hack` prefix will be used")
+
+            val renames =
+              model
+                .shapes()
+                .collect(Collectors.toList())
+                .asScala
+                .filter(_.getId().getNamespace() == args.smithyNamespace)
+                .map { shp =>
+                  shp.getId() ->
+                    shp.getId().withNamespace(renameNamespaceForHack(shp.getId().getNamespace()))
+                }
+                .toMap
+                .asJava
+
+            ModelTransformer
+              .create()
+              .renameShapes(model, renames)
+          } else
+            model
       }
-      .assemble()
-      .unwrap()
+
     val context = PluginContext
       .builder()
       .model(model)
@@ -128,15 +177,35 @@ object SmithyTraitCodegen {
         ObjectNode
           .builder()
           .withMember("package", args.javaPackage)
-          .withMember("namespace", args.smithyNamespace)
+          .withMember("namespace", renameNamespaceForHack(args.smithyNamespace))
           .withMember("header", ArrayNode.builder.build())
           .withMember("excludeTags", ArrayNode.builder.withValue("nocodegen").build())
           .build()
       )
       .build()
     val plugin = new TraitCodegenPlugin()
     plugin.execute(context)
-    os.move(genDir / "META-INF", metaDir / "META-INF")
+
+    // If there were no shapes to generate, this won't exist
+    if (os.exists(genDir / "META-INF"))
+      os.move(genDir / "META-INF", metaDir / "META-INF")
+
+    os.walk(genDir)
+      .filter(os.isFile)
+      .filter(_.ext == "java")
+      .foreach { f =>
+        os.write.over(f, replaceNamespaceRefsInFile(os.read(f), args.smithyNamespace))
+      }
+
+    os
+      .walk(metaDir, includeTarget = true)
+      .filter(os.isFile)
+      .foreach { p =>
+        if (p.toIO.name == "software.amazon.smithy.model.traits.TraitService") {
+          args.externalProviders.foreach(provider => os.write.append(p, provider + "\n"))
+        }
+      }
+
     Output(metaDir = metaDir.toIO, javaDir = genDir.toIO)
   }
 

sbtPlugin/src/main/scala/org/polyvariant/smithytraitcodegen/SmithyTraitCodegenPlugin.scala

@@ -39,11 +39,26 @@ object SmithyTraitCodegenPlugin extends AutoPlugin {
       "Dependencies to be added into codegen model"
     )
 
+    val smithyTraitCodegenSourceDirectory = settingKey[File](
+      "The directory where the smithy sources are located"
+    )
+
+    val smithyTraitCodegenTargetDirectory = settingKey[File](
+      "The directory where the generated Java sources and resources will be placed"
+    )
+
+    val smithyTraitCodegenExternalProviders = settingKey[List[String]](
+      "External trait provideres"
+    )
+
   }
 
   import autoImport.*
 
   override def projectSettings: Seq[Setting[?]] = Seq(
+    smithyTraitCodegenSourceDirectory := (Compile / resourceDirectory).value / "META-INF" / "smithy",
+    smithyTraitCodegenTargetDirectory := (Compile / target).value,
+    smithyTraitCodegenExternalProviders := Nil,
     Keys.generateSmithyTraits := Def.task {
       import sbt.util.CacheImplicits.*
       val s = (Compile / streams).value
@@ -64,9 +79,10 @@ object SmithyTraitCodegenPlugin extends AutoPlugin {
       val args = SmithyTraitCodegen.Args(
         javaPackage = smithyTraitCodegenJavaPackage.value,
         smithyNamespace = smithyTraitCodegenNamespace.value,
-        targetDir = os.Path((Compile / target).value),
-        smithySourcesDir = PathRef((Compile / resourceDirectory).value / "META-INF" / "smithy"),
+        targetDir = os.Path(smithyTraitCodegenTargetDirectory.value),
+        smithySourcesDir = PathRef(smithyTraitCodegenSourceDirectory.value),
         dependencies = jars.map(PathRef(_)).toList,
+        externalProviders = smithyTraitCodegenExternalProviders.value,
       )
       val cachedCodegen =
         Tracked.inputChanged[SmithyTraitCodegen.Args, SmithyTraitCodegen.Output](
@@ -82,7 +98,7 @@ object SmithyTraitCodegenPlugin extends AutoPlugin {
                   .fold {
                     SmithyTraitCodegen.generate(codegenArgs)
                   } { last =>
-                    logger.info(s"Using cached result of smithy-trait-codegen")
+                    logger.info("Using cached result of smithy-trait-codegen")
                     last
                   }
               }

sbtPlugin/src/sbt-test/codegen/namespace-restriction-hack/build.sbt

@@ -0,0 +1,13 @@
+val traits = project
+  .enablePlugins(SmithyTraitCodegenPlugin)
+  .settings(
+    smithyTraitCodegenDependencies := List(
+      "io.github.disneystreaming.alloy" % "alloy-core" % "0.3.23"
+    ),
+    smithyTraitCodegenJavaPackage := "smithy4bazinga",
+    smithyTraitCodegenNamespace := "smithy4bazinga",
+  )
+
+val root = project
+  .in(file("."))
+  .dependsOn(traits)

sbtPlugin/src/sbt-test/codegen/namespace-restriction-hack/project/plugins.sbt

@@ -0,0 +1,5 @@
+sys.props.get("plugin.version") match {
+  case Some(x) => addSbtPlugin("org.polyvariant" % "smithy-trait-codegen-sbt" % x)
+  case _       => sys.error("""|The system property 'plugin.version' is not defined.
+                         |Specify this property using the scriptedLaunchOpts -D.""".stripMargin)
+}

sbtPlugin/src/sbt-test/codegen/namespace-restriction-hack/src/main/scala/Demo.scala

@@ -0,0 +1,8 @@
+import software.amazon.smithy.model.Model
+
+object DemoApp extends App {
+  assert(
+    smithy4bazinga.BazingaTrait.ID.toString() == "smithy4bazinga#bazinga",
+    s"Expected smithy4bazinga.BazingaTrait.ID to be 'smithy4bazinga#bazinga', but got ${smithy4bazinga.BazingaTrait.ID.toString()}",
+  )
+}

sbtPlugin/src/sbt-test/codegen/namespace-restriction-hack/test

@@ -0,0 +1 @@
+> run

sbtPlugin/src/sbt-test/codegen/namespace-restriction-hack/traits/src/main/resources/META-INF/smithy/manifest

@@ -0,0 +1 @@
+smithy4bazinga.smithy

sbtPlugin/src/sbt-test/codegen/namespace-restriction-hack/traits/src/main/resources/META-INF/smithy/smithy4bazinga.smithy

@@ -0,0 +1,6 @@
+$version: "2"
+
+namespace smithy4bazinga
+
+@trait
+string bazinga