Merge pull request #273 from mpollmeier/project-by

type safe `project` step by To-om

Author: mpollmeier

gremlin-scala/src/main/scala/gremlin/scala/GremlinScala.scala

@@ -33,6 +33,7 @@ import org.apache.tinkerpop.gremlin.structure.{Direction, T}
 import shapeless.{::, HList, HNil}
 import shapeless.ops.hlist.{IsHCons, Mapper, Prepend, RightFolder, ToTraversable, Tupler}
 import shapeless.ops.product.ToHList
+import shapeless.syntax.std.tuple._
 import scala.concurrent.duration.FiniteDuration
 import scala.reflect.runtime.{universe => ru}
 import scala.collection.{immutable, mutable}
@@ -116,6 +117,10 @@ class GremlinScala[End](val traversal: GraphTraversal[_, End]) {
                  otherProjectKeys: String*): GremlinScala.Aux[JMap[String, A], Labels] =
     GremlinScala[JMap[String, A], Labels](traversal.project(projectKey, otherProjectKeys: _*))
 
+  def project[H <: Product](
+      builder: ProjectionBuilder[Nil.type] => ProjectionBuilder[H]): GremlinScala[H] =
+    builder(ProjectionBuilder()).build(this)
+
   /** You might think that predicate should be `GremlinScala[End] => GremlinScala[Boolean]`,
     * but that's not how tp3 works: e.g. `.value(Age).is(30)` returns `30`, not `true`
     */

gremlin-scala/src/main/scala/gremlin/scala/ProjectionBuilder.scala

@@ -0,0 +1,31 @@
+package gremlin.scala
+
+import java.util.{Map => JMap, UUID}
+import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversal
+import shapeless.ops.tuple.{Prepend => TuplePrepend}
+import shapeless.syntax.std.tuple._
+
+class ProjectionBuilder[T <: Product] private[gremlin] (
+    labels: Seq[String],
+    addBy: GraphTraversal[_, JMap[String, Any]] => GraphTraversal[_, JMap[String, Any]],
+    buildResult: JMap[String, Any] => T) {
+
+  def apply[U, TR <: Product](by: By[U])(
+      implicit prepend: TuplePrepend.Aux[T, Tuple1[U], TR]): ProjectionBuilder[TR] = {
+    val label = UUID.randomUUID().toString
+    new ProjectionBuilder[TR](labels :+ label,
+                              addBy.andThen(by.apply),
+                              map => buildResult(map) :+ map.get(label).asInstanceOf[U])
+  }
+
+  def and[U, TR <: Product](by: By[U])(
+      implicit prepend: TuplePrepend.Aux[T, Tuple1[U], TR]): ProjectionBuilder[TR] = apply(by)
+
+  private[gremlin] def build(g: GremlinScala[_]): GremlinScala[T] = {
+    GremlinScala(addBy(g.traversal.project(labels.head, labels.tail: _*))).map(buildResult)
+  }
+}
+
+object ProjectionBuilder {
+  def apply() = new ProjectionBuilder[Nil.type](Nil, scala.Predef.identity, _ => Nil)
+}

gremlin-scala/src/test/scala/gremlin/scala/ProjectSpec.scala

@@ -0,0 +1,36 @@
+package gremlin.scala
+
+import org.apache.tinkerpop.gremlin.tinkergraph.structure.TinkerFactory
+import org.scalatest.{Matchers, WordSpec}
+
+class ProjectSpec extends WordSpec with Matchers {
+
+  "projecting by two traversals" in {
+    val result: (java.lang.Long, java.lang.Long) =
+      graph.V
+        .has(name.of("marko"))
+        .project(_(By(__.outE.count)).and(By(__.inE.count)))
+        .head
+
+    result shouldBe (3, 0)
+  }
+
+  "projecting by property and traversal" in {
+    val result: List[(String, java.lang.Long)] =
+      graph.V
+        .out("created")
+        .project(_(By(name)).and(By(__.in("created").count)))
+        .toList
+
+    result shouldBe List(
+      ("lop", 3),
+      ("lop", 3),
+      ("lop", 3),
+      ("ripple", 1)
+    )
+  }
+
+  def graph: ScalaGraph = TinkerFactory.createModern.asScala
+  val name = Key[String]("name")
+
+}