Add companion_extends

Fix scalapb#105.
trungfinity · Apr 28, 2016 · b52e9fd · b52e9fd
1 parent b9de661
commit b52e9fd
Show file tree

Hide file tree

Showing 8 changed files with 582 additions and 61 deletions.
diff --git a/compiler-plugin/src/main/java/com/trueaccord/scalapb/Scalapb.java b/compiler-plugin/src/main/java/com/trueaccord/scalapb/Scalapb.java
diff --git a/compiler-plugin/src/main/scala/com/trueaccord/scalapb/compiler/DescriptorPimps.scala b/compiler-plugin/src/main/scala/com/trueaccord/scalapb/compiler/DescriptorPimps.scala
@@ -226,13 +226,23 @@ trait DescriptorPimps {
 
     def extendsOption = messageOptions.getExtendsList.toSeq
 
+    def companionExtendsOption = messageOptions.getCompanionExtendsList.toSeq
+
     def nameSymbol = scalaName.asSymbol
 
     def baseClasses: Seq[String] =
       Seq("com.trueaccord.scalapb.GeneratedMessage",
         s"com.trueaccord.scalapb.Message[$nameSymbol]",
         s"com.trueaccord.lenses.Updatable[$nameSymbol]") ++ extendsOption
 
+    def companionBaseClasses: Seq[String] = {
+      val mixins = if (javaConversions)
+        Seq(s"com.trueaccord.scalapb.JavaProtoSupport[$nameSymbol, $javaTypeName]") else Nil
+      Seq(s"com.trueaccord.scalapb.GeneratedMessageCompanion[$nameSymbol]") ++
+        mixins ++
+        companionExtendsOption
+    }
+
     def nestedTypes: Seq[Descriptor] = message.getNestedTypes.toSeq
 
     def isMapEntry: Boolean = message.getOptions.getMapEntry

diff --git a/compiler-plugin/src/main/scala/com/trueaccord/scalapb/compiler/ProtobufGenerator.scala b/compiler-plugin/src/main/scala/com/trueaccord/scalapb/compiler/ProtobufGenerator.scala
@@ -809,9 +809,7 @@ class ProtobufGenerator(val params: GeneratorParams) extends DescriptorPimps {
 
   def generateMessageCompanion(message: Descriptor)(printer: FunctionalPrinter): FunctionalPrinter = {
     val className = message.nameSymbol
-    val mixins = if (message.javaConversions)
-      s"with com.trueaccord.scalapb.JavaProtoSupport[$className, ${message.javaTypeName}] " else ""
-    val companionType = s"com.trueaccord.scalapb.GeneratedMessageCompanion[$className] $mixins"
+    val companionType = message.companionBaseClasses.mkString(" with ")
     printer.addM(
       s"""object $className extends $companionType {
          |  implicit def messageCompanion: $companionType = this""")
@@ -975,7 +973,7 @@ class ProtobufGenerator(val params: GeneratorParams) extends DescriptorPimps {
   }
 
   def generateSingleScalaFileForFileDescriptor(file: FileDescriptor): Seq[CodeGeneratorResponse.File] = {
-    val code = 
+    val code =
       scalaFileHeader(file)
       .print(file.getEnumTypes)(printEnum)
       .print(file.getMessageTypes)(printMessage)
@@ -1073,8 +1071,8 @@ object ProtobufGenerator {
           request.getFileToGenerateList.foreach {
             name =>
               val file = filesByName(name)
-              val responseFiles = 
-                if (file.scalaOptions.getSingleFile) 
+              val responseFiles =
+                if (file.scalaOptions.getSingleFile)
                   generator.generateSingleScalaFileForFileDescriptor(file)
                 else generator.generateMultipleScalaFilesForFileDescriptor(file)
               b.addAllFile(responseFiles)

diff --git a/e2e/src/main/protobuf/custom_types.proto b/e2e/src/main/protobuf/custom_types.proto
@@ -47,6 +47,7 @@ message OneofMessage {
 
 message CustomerEvent {
   option (scalapb.message).extends = "com.trueaccord.pb.DomainEvent";
+  option (scalapb.message).companion_extends = "com.trueaccord.pb.DomainEventCompanion";
 
   optional string person_id = 1          [(scalapb.field).type = "com.trueaccord.pb.PersonId"];
   optional int32 optional_number = 2;

diff --git a/e2e/src/main/scala/com/trueaccord/pb/PersonId.scala b/e2e/src/main/scala/com/trueaccord/pb/PersonId.scala
@@ -36,3 +36,7 @@ trait DomainEvent {
   def repeatedNumber: Seq[Int]
   def requiredNumber: Int
 }
+
+trait DomainEventCompanion {
+  val thisIs = "The companion object"
+}
diff --git a/e2e/src/test/scala/CustomTypesSpec.scala b/e2e/src/test/scala/CustomTypesSpec.scala
@@ -98,4 +98,9 @@ class CustomTypesSpec extends FlatSpec with MustMatchers {
     t.repeatedNumber must be(Seq(2,3,4))
     t.requiredNumber must be(5)
   }
+
+  "Extended companion objects" should "inherit from marker type" in {
+    CustomerEvent mustBe a [DomainEventCompanion]
+    CustomerEvent.thisIs must be("The companion object")
+  }
 }
diff --git a/protobuf/scalapb/scalapb.proto b/protobuf/scalapb/scalapb.proto
@@ -36,8 +36,11 @@ extend google.protobuf.FileOptions {
 }
 
 message MessageOptions {
-  // additional classes and traits to mix in to the case class.
+  // Additional classes and traits to mix in to the case class.
   repeated string extends = 1;
+
+  // Additional classes and traits to mix in to the companion object.
+  repeated string companion_extends = 2;
 }
 
 extend google.protobuf.MessageOptions {