move packageCode override to separate trait

2 files changed, 29 insertions, 30 deletions

diff --git a/src/main/scala/NamespacedCodegen.scala b/src/main/scala/NamespacedCodegen.scala
index 52761b9..affd36c 100644
--- a/src/main/scala/NamespacedCodegen.scala
+++ b/src/main/scala/NamespacedCodegen.scala
@@ -6,7 +6,6 @@ import scala.concurrent.duration.Duration
 import scala.concurrent.ExecutionContext.Implicits.global
 import slick.backend.DatabaseConfig
 import slick.codegen.{
-  OutputHelpers,
   SourceCodeGenerator,
   StringGeneratorHelpers
 }
@@ -70,12 +69,6 @@ object Generator {
 
 }
 
-class ImportGenerator(dbModel: Model, schemaImports: List[String])
-    extends SourceCodeGenerator(dbModel) {
-  override def code: String =
-    schemaImports.map("import " + _).mkString("\n") + "\n"
-}
-
 class Generator(pkg: String,
                 fullDatabaseModel: Model,
                 schemaOnlyModel: Model,
@@ -87,8 +80,7 @@ class Generator(pkg: String,
     extends SourceCodeGenerator(schemaOnlyModel)
     with OutputHelpers {
 
-  val allImports: String =
-    new ImportGenerator(fullDatabaseModel, schemaImports).code
+  override val imports = schemaImports.map("import " + _).mkString("\n")
 
   val defaultIdImplementation =
     """|final case class Id[T](v: Int)
@@ -104,27 +96,6 @@ class Generator(pkg: String,
   // Alias to ForeignKeyAction is in profile.api
   // TODO: fix upstream
 
-  override def packageCode(profile: String,
-                           pkg: String,
-                           container: String,
-                           parentType: Option[String]): String = {
-    val traitName = container.capitalize + "SchemaDef"
-    s"""|package $pkg
-        |$allImports
-        |// AUTO-GENERATED Slick data model
-        |
-        |/** Stand-alone Slick data model for immediate use */
-        |object $container extends {
-        |  val profile = $profile
-        |} with $traitName
-        |
-        |/** Slick data model trait for extension, choice of backend or usage in the cake pattern. (Make sure to initialize this late.) */
-        |trait $traitName${parentType.fold("")(" extends " + _)} {
-        |  import profile.api._
-        |  ${indent(code)}
-        |}""".stripMargin.trim()
-  }
-
   override def Table = new Table(_) { table =>
 
     // need this in order to use our own TableClass generator
diff --git a/src/main/scala/OutputHelpers.scala b/src/main/scala/OutputHelpers.scala
new file mode 100644
index 0000000..f982124
--- /dev/null
+++ b/src/main/scala/OutputHelpers.scala
@@ -0,0 +1,28 @@
+trait OutputHelpers extends slick.codegen.OutputHelpers {
+
+  def imports: String
+
+  def headerComment: String = ""
+
+  override def packageCode(profile: String,
+                           pkg: String,
+                           container: String,
+                           parentType: Option[String]): String = {
+    val traitName = container.capitalize + "SchemaDef"
+    s"""|package $pkg
+        |$imports
+        |
+        |// AUTO-GENERATED Slick data model
+        |
+        |/** Stand-alone Slick data model for immediate use */
+        |object $container extends {
+        |  val profile = $profile
+        |} with $traitName
+        |
+        |/** Slick data model trait for extension, choice of backend or usage in the cake pattern. (Make sure to initialize this late.) */
+        |trait $traitName${parentType.fold("")(" extends " + _)} {
+        |  import profile.api._
+        |  ${indent(code)}
+        |}""".stripMargin.trim()
+  }
+}