Allow to generate more convenient setters for option fields (#19)

Use like
```scala
@DaTa(optionSetters = true) class Foo(n: Option[Int] = None)
Foo().withN(2) // wrapping in Some(…) is optional
```

Author: alexarchambault

src/main/scala/dataclass/Macros.scala

@@ -13,8 +13,10 @@ private[dataclass] class Macros(val c: Context) extends ImplTransformers {
     .map(_.toBoolean)
     .getOrElse(java.lang.Boolean.getBoolean("dataclass.macros.debug"))
 
-  private class Transformer(generateApplyMethods: Boolean)
-      extends ImplTransformer {
+  private class Transformer(
+      generateApplyMethods: Boolean,
+      generateOptionSetters: Boolean
+  ) extends ImplTransformer {
     override def transformClass(
         cdef: ClassDef,
         mdef: ModuleDef
@@ -92,7 +94,7 @@ private[dataclass] class Macros(val c: Context) extends ImplTransformers {
 
           val setters = paramss.zipWithIndex.flatMap {
             case (l, groupIdx) =>
-              l.zipWithIndex.map {
+              l.zipWithIndex.flatMap {
                 case (p, idx) =>
                   val namedArgs0 =
                     namedArgs.updated(
@@ -101,7 +103,34 @@ private[dataclass] class Macros(val c: Context) extends ImplTransformers {
                     )
                   val fn = p.name.decodedName.toString.capitalize
                   val withDefIdent = TermName(s"with$fn")
-                  q"def $withDefIdent(${p.name}: ${p.tpt}) = new $tpname[..$tparamsRef](...$namedArgs0)"
+
+                  val extraMethods =
+                    if (generateOptionSetters) {
+                      val wrappedOptionTpe = p.tpt match {
+                        case AppliedTypeTree(
+                            Ident(TypeName("Option")),
+                            List(wrapped)
+                            ) =>
+                          Seq(wrapped)
+                        case _ => Nil
+                      }
+
+                      wrappedOptionTpe.map { tpe0 =>
+                        val namedArgs0 =
+                          namedArgs.updated(
+                            groupIdx,
+                            namedArgs(groupIdx).updated(
+                              idx,
+                              q"${p.name}=_root_.scala.Some(${p.name})"
+                            )
+                          )
+                        q"def $withDefIdent(${p.name}: $tpe0) = new $tpname[..$tparamsRef](...$namedArgs0)"
+                      }
+                    } else
+                      Nil
+
+                  q"def $withDefIdent(${p.name}: ${p.tpt}) = new $tpname[..$tparamsRef](...$namedArgs0)" +:
+                    extraMethods
               }
           }
 
@@ -422,7 +451,14 @@ private[dataclass] class Macros(val c: Context) extends ImplTransformers {
       case _              => true
     }
 
-    annottees.transformAnnottees(new Transformer(generateApplyMethods))
+    val generateOptionSetters = params.exists {
+      case q"optionSetters=true" => true
+      case _                     => false
+    }
+
+    annottees.transformAnnottees(
+      new Transformer(generateApplyMethods, generateOptionSetters)
+    )
   }
 
 }

src/main/scala/dataclass/data.scala

@@ -4,7 +4,11 @@ import scala.annotation.{StaticAnnotation, compileTimeOnly}
 import scala.language.experimental.macros
 
 @compileTimeOnly("enable macro paradise to expand macro annotations")
-class data(apply: Boolean = true, publicConstructor: Boolean = true)
-    extends StaticAnnotation {
+class data(
+    apply: Boolean = true,
+    publicConstructor: Boolean = true,
+    /** Whether to generate `withFoo(foo: Foo)` methods for fields like `foo: Option[Foo]`) */
+    optionSetters: Boolean = false
+) extends StaticAnnotation {
   def macroTransform(annottees: Any*): Any = macro Macros.impl
 }

src/test/scala/dataclass/MoreFieldsTests.scala

@@ -20,6 +20,23 @@ object MoreFieldsTests extends TestSuite {
       assert(t == (1, "a", false, 1.2))
     }
 
+    "option setters" - {
+      @data(optionSetters = true) class Bar(
+          n: Int,
+          s: String,
+          b: Option[Boolean] = None,
+          d: Double = 1.0
+      )
+      val foo = Bar(1, "a")
+      val foo2 = foo.withB(true)
+      val foo3 = foo.withB(Some(false))
+      val foo4 = foo2.withB(None)
+      assert(foo.b == None)
+      assert(foo2.b == Some(true))
+      assert(foo3.b == Some(false))
+      assert(foo4.b == None)
+    }
+
     "since annotation" - {
       * - {
         illTyped(

src/test/scala/dataclass/OneFieldTests.scala

@@ -49,6 +49,21 @@ object OneFieldTests extends TestSuite {
       assert(foo.count == 1)
       assert(foo2.count == 2)
     }
+    "option setter" - {
+      @data(optionSetters = true) class Bar(count: Option[Int] = None)
+      val bar = Bar()
+      val bar2 = bar.withCount(2)
+      assert(bar.count == None)
+      assert(bar2.count == Some(2))
+    }
+    "no option setter" - {
+      @data class Bar(count: Option[Int] = None)
+      val bar = Bar()
+      illTyped("""
+        val bar2 = bar.withCount(2)
+      """, "type mismatch;.*")
+      assert(bar.count == None)
+    }
 
     "tuple" - {
       @data class Foo0(a: Int) {

src/test/scala/dataclass/TwoFieldsTests.scala

@@ -62,6 +62,27 @@ object TwoFieldsTests extends TestSuite {
       assert(foo4.a == 2)
       assert(foo4.other == "u")
     }
+    "option setters" - {
+      @data(optionSetters = true) class Bar(
+          a: Int,
+          other: Option[String] = None
+      )
+      val bar = Bar(1)
+      val bar2 = bar.withA(2)
+      val bar3 = bar.withOther("t")
+      val bar4 = bar2.withOther(Some("u"))
+      val bar5 = bar2.withOther(None)
+      assert(bar.a == 1)
+      assert(bar.other == None)
+      assert(bar2.a == 2)
+      assert(bar2.other == None)
+      assert(bar3.a == 1)
+      assert(bar3.other == Some("t"))
+      assert(bar4.a == 2)
+      assert(bar4.other == Some("u"))
+      assert(bar5.a == 2)
+      assert(bar5.other == None)
+    }
 
     "tuple" - {
       @data class Foo0(a: Int, other: String) {