diff --git a/gremlin-scala/src/main/scala/gremlin/scala/GremlinScala.scala b/gremlin-scala/src/main/scala/gremlin/scala/GremlinScala.scala index f35f9ec1..55cfa9c2 100644 --- a/gremlin-scala/src/main/scala/gremlin/scala/GremlinScala.scala +++ b/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` */ diff --git a/gremlin-scala/src/main/scala/gremlin/scala/ProjectionBuilder.scala b/gremlin-scala/src/main/scala/gremlin/scala/ProjectionBuilder.scala new file mode 100644 index 00000000..700d5a0a --- /dev/null +++ b/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) +} 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..9713e5a0 --- /dev/null +++ b/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") + +}