diff --git a/gremlin-scala/src/main/scala/gremlin/scala/GremlinScala.scala b/gremlin-scala/src/main/scala/gremlin/scala/GremlinScala.scala index f35f9ec1..39cd7b2c 100644 --- a/gremlin-scala/src/main/scala/gremlin/scala/GremlinScala.scala +++ b/gremlin-scala/src/main/scala/gremlin/scala/GremlinScala.scala @@ -13,7 +13,8 @@ import java.util.{ Map => JMap, Collection => JCollection, Iterator => JIterator, - Set => JSet + Set => JSet, + UUID } import java.util.stream.{Stream => JStream} @@ -32,7 +33,9 @@ import org.apache.tinkerpop.gremlin.process.traversal.{Bytecode, Path, Scope, Tr 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.tuple.{Prepend => TuplePrepend} 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 +119,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` */ @@ -1027,3 +1034,28 @@ class GremlinScala[End](val traversal: GraphTraversal[_, End]) { travs.map(_.apply(start).traversal) } + +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) +} diff --git a/gremlin-scala/src/test/scala/gremlin/scala/ProjectSpec.scala b/gremlin-scala/src/test/scala/gremlin/scala/ProjectSpec.scala new file mode 100644 index 00000000..a47e72f7 --- /dev/null +++ b/gremlin-scala/src/test/scala/gremlin/scala/ProjectSpec.scala @@ -0,0 +1,37 @@ +package gremlin.scala + +import org.apache.tinkerpop.gremlin.tinkergraph.structure.TinkerFactory +import org.scalatest.{Matchers, WordSpec} + +class ProjectSpec extends WordSpec with Matchers { + def graph: ScalaGraph = TinkerFactory.createModern.asScala() + + "project steps" should { + "provide type safe result" in { + val result = graph + .V() + .out("created") + .project(_(By(Key[String]("name"))) + .and(By(__.in("created").count()))) + .toList() + + result shouldBe List( + ("lop", 3), + ("lop",3), + ("lop",3), + ("ripple", 1) + ) + } + + "provide other type safe result" in { + val result = graph + .V() + .has(Key("name").of("marko")) + .project(_(By(__.outE().count())) + .and(By(__.inE().count()))) + .head() + + result shouldBe (3, 0) + } + } +}