From 81bbeb30195b5d991fd9bcbb2dd22c5149397537 Mon Sep 17 00:00:00 2001
From: Nadav Samet <thesamet@gmail.com>
Date: Sat, 30 Jan 2021 09:28:39 -0800
Subject: [PATCH] Allow disabling flat_package from file options

See scalapb/scalapb-validate#70
---
 .../compiler/DescriptorImplicits.scala        |  13 ++-
 .../compiler/DescriptorImplicitsSpec.scala    | 101 ++++++++++++++++++
 2 files changed, 111 insertions(+), 3 deletions(-)
 create mode 100644 compiler-plugin/src/test/scala/scalapb/compiler/DescriptorImplicitsSpec.scala

diff --git a/compiler-plugin/src/main/scala/scalapb/compiler/DescriptorImplicits.scala b/compiler-plugin/src/main/scala/scalapb/compiler/DescriptorImplicits.scala
index b13c0dfb2..b1ef5c188 100644
--- a/compiler-plugin/src/main/scala/scalapb/compiler/DescriptorImplicits.scala
+++ b/compiler-plugin/src/main/scala/scalapb/compiler/DescriptorImplicits.scala
@@ -906,15 +906,22 @@ class DescriptorImplicits private[compiler] (
         else r + "OuterClass"
       }
 
-    private def isNonFlatDependency =
-      (file.getPackage == "google.protobuf") || (file.getPackage == "scalapb")
+    private def isFlatPackage = {
+      // Disable flat-package for these two packages, even if a generator
+      // option is given:
+      val isFlatPackageSpecialCase =
+        (file.getPackage == "google.protobuf") || (file.getPackage == "scalapb")
+
+      if (scalaOptions.hasFlatPackage()) scalaOptions.getFlatPackage()
+      else (params.flatPackage && !isFlatPackageSpecialCase)
+    }
 
     private def scalaPackageParts: Seq[String] = {
       val requestedPackageName: Seq[String] =
         (if (scalaOptions.hasPackageName) scalaOptions.getPackageName.split('.')
          else javaPackage.split('.')).toIndexedSeq.filterNot(_.isEmpty)
 
-      if (scalaOptions.getFlatPackage || (params.flatPackage && !isNonFlatDependency))
+      if (isFlatPackage)
         requestedPackageName
       else requestedPackageName ++ baseName(file.getName).replace('-', '_').split('.')
     }
diff --git a/compiler-plugin/src/test/scala/scalapb/compiler/DescriptorImplicitsSpec.scala b/compiler-plugin/src/test/scala/scalapb/compiler/DescriptorImplicitsSpec.scala
new file mode 100644
index 000000000..97cddc96c
--- /dev/null
+++ b/compiler-plugin/src/test/scala/scalapb/compiler/DescriptorImplicitsSpec.scala
@@ -0,0 +1,101 @@
+package scalapb.compiler
+
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.must.Matchers
+
+class DescriptorImplicitsSpec extends AnyFlatSpec with Matchers with ProtocInvocationHelper {
+  val base = Seq(
+    "disable_flat.proto" ->
+      """|syntax = "proto2";
+         |package disable_flat;
+         |import "scalapb/scalapb.proto";
+         |option (scalapb.options) = {
+         |  flat_package: false;
+         |  scope: PACKAGE
+         |};
+         |""".stripMargin,
+    "enable_flat.proto" ->
+      """|syntax = "proto2";
+         |package enable_flat;
+         |import "scalapb/scalapb.proto";
+         |option (scalapb.options) = {
+         |  flat_package: true;
+         |  scope: PACKAGE
+         |};
+         |message Foo {}
+         |""".stripMargin,
+    "inside_disable_flat.proto" ->
+      """|syntax = "proto2";
+         |package disable_flat;
+         |message A {};
+         |""".stripMargin,
+    "inside_enable_flat.proto" ->
+      """|syntax = "proto2";
+         |package enable_flat;
+         |message B {};
+         |""".stripMargin,
+    "outside.proto" ->
+      """|syntax = "proto2";
+         |package outside;
+         |message C {};
+         |""".stripMargin
+  )
+
+  "flat package" should "be overridable to false when set as generator parameter" in {
+    val files = generateFileSet(base)
+    val implicits = new DescriptorImplicits(
+      GeneratorParams(flatPackage = true),
+      files,
+      SecondaryOutputProvider.empty
+    )
+    import implicits._
+
+    files
+      .find(_.getFullName() == "inside_disable_flat.proto")
+      .get
+      .findMessageTypeByName("A")
+      .scalaType
+      .fullName must be("disable_flat.inside_disable_flat.A")
+    files
+      .find(_.getFullName() == "inside_enable_flat.proto")
+      .get
+      .findMessageTypeByName("B")
+      .scalaType
+      .fullName must be("enable_flat.B")
+    files
+      .find(_.getFullName() == "outside.proto")
+      .get
+      .findMessageTypeByName("C")
+      .scalaType
+      .fullName must be("outside.C")
+  }
+
+  "flat package" should "be overridable when not set as generator parameter" in {
+    val files = generateFileSet(base)
+    val implicits = new DescriptorImplicits(
+      GeneratorParams(flatPackage = false),
+      files,
+      SecondaryOutputProvider.empty
+    )
+    import implicits._
+
+    files
+      .find(_.getFullName() == "inside_disable_flat.proto")
+      .get
+      .findMessageTypeByName("A")
+      .scalaType
+      .fullName must be("disable_flat.inside_disable_flat.A")
+    files
+      .find(_.getFullName() == "inside_enable_flat.proto")
+      .get
+      .findMessageTypeByName("B")
+      .scalaType
+      .fullName must be("enable_flat.B")
+    files
+      .find(_.getFullName() == "outside.proto")
+      .get
+      .findMessageTypeByName("C")
+      .scalaType
+      .fullName must be("outside.outside.C")
+  }
+}