From 173993a797814cc0ebba76144c527965d659e394 Mon Sep 17 00:00:00 2001
From: Michael Pollmeier <michael@michaelpollmeier.com>
Date: Wed, 26 Jun 2024 12:47:02 +0200
Subject: [PATCH] NumericSteps etc.: some more functionality and tests from odb
 (#210)

---
 .../scala/flatgraph/traversal/Language.scala  |  3 ++
 .../flatgraph/traversal/NumericSteps.scala    | 31 +++++++++++++++++++
 .../traversal/GratefulDeadTests.scala         |  7 +++++
 .../flatgraph/traversal/TraversalTests.scala  | 14 +++++++--
 4 files changed, 52 insertions(+), 3 deletions(-)
 create mode 100644 core/src/main/scala/flatgraph/traversal/NumericSteps.scala

diff --git a/core/src/main/scala/flatgraph/traversal/Language.scala b/core/src/main/scala/flatgraph/traversal/Language.scala
index 9f2dae40..8e53edee 100644
--- a/core/src/main/scala/flatgraph/traversal/Language.scala
+++ b/core/src/main/scala/flatgraph/traversal/Language.scala
@@ -26,6 +26,9 @@ trait language {
 
   implicit def iteratorToEdgeSteps[A <: Edge](iter: IterableOnce[A]): EdgeSteps[A] =
     new EdgeSteps[A](iter.iterator)
+
+  implicit def iteratorToNumericSteps[A: Numeric](iter: IterableOnce[A]): NumericSteps[A] =
+    new NumericSteps[A](iter)
 }
 
 @Traversal(elementType = classOf[AnyRef])
diff --git a/core/src/main/scala/flatgraph/traversal/NumericSteps.scala b/core/src/main/scala/flatgraph/traversal/NumericSteps.scala
new file mode 100644
index 00000000..ff0e3b53
--- /dev/null
+++ b/core/src/main/scala/flatgraph/traversal/NumericSteps.scala
@@ -0,0 +1,31 @@
+package flatgraph.traversal
+
+class NumericSteps[N](val traversal: Iterator[N]) extends AnyVal {
+
+  def greaterThan(n: N)(implicit numeric: Numeric[N]): Iterator[N] =
+    traversal.filter(numeric.gt(_, n))
+
+  def greaterThanEqual(n: N)(implicit numeric: Numeric[N]): Iterator[N] =
+    traversal.filter(numeric.gteq(_, n))
+
+  def lessThan(n: N)(implicit numeric: Numeric[N]): Iterator[N] =
+    traversal.filter(numeric.lt(_, n))
+
+  def lessThanEqual(n: N)(implicit numeric: Numeric[N]): Iterator[N] =
+    traversal.filter(numeric.lteq(_, n))
+
+  def equal(n: N)(implicit numeric: Numeric[N]): Iterator[N] =
+    traversal.filter(numeric.equiv(_, n))
+
+  def equiv(n: N)(implicit numeric: Numeric[N]): Iterator[N] = equal(n)
+
+  def between(startInclusive: N, endExclusive: N)(implicit numeric: Numeric[N]): Iterator[N] =
+    traversal.filter(n => numeric.gteq(n, startInclusive) && numeric.lt(n, endExclusive))
+
+  def inside(startExclusive: N, endExclusive: N)(implicit numeric: Numeric[N]): Iterator[N] =
+    traversal.filter(n => numeric.gt(n, startExclusive) && numeric.lt(n, endExclusive))
+
+  def outside(startInclusive: N, endInclusive: N)(implicit numeric: Numeric[N]): Iterator[N] =
+    traversal.filter(n => numeric.lt(n, startInclusive) || numeric.gt(n, endInclusive))
+
+}
diff --git a/tests/src/test/scala/flatgraph/traversal/GratefulDeadTests.scala b/tests/src/test/scala/flatgraph/traversal/GratefulDeadTests.scala
index 6ab03d10..f9465e1e 100644
--- a/tests/src/test/scala/flatgraph/traversal/GratefulDeadTests.scala
+++ b/tests/src/test/scala/flatgraph/traversal/GratefulDeadTests.scala
@@ -77,6 +77,13 @@ class GratefulDeadTests extends AnyWordSpec {
       gratefulDead.song.performancesGte(1).size shouldBe 483
       gratefulDead.song.performancesLt(1).size shouldBe 101
       gratefulDead.song.performancesLte(1).size shouldBe 243
+
+      // numeric filter steps
+      gratefulDead.song.where(_.performances.equal(1)).size shouldBe 142
+      gratefulDead.song.where(_.performances.greaterThan(1)).size shouldBe 341
+      gratefulDead.song.where(_.performances.greaterThanEqual(1)).size shouldBe 483
+      gratefulDead.song.where(_.performances.lessThan(1)).size shouldBe 101
+      gratefulDead.song.where(_.performances.lessThanEqual(1)).size shouldBe 243
     }
 
     "throw useful exception when passing invalid regexp" in {
diff --git a/tests/src/test/scala/flatgraph/traversal/TraversalTests.scala b/tests/src/test/scala/flatgraph/traversal/TraversalTests.scala
index 997a0d4b..bd22ade7 100644
--- a/tests/src/test/scala/flatgraph/traversal/TraversalTests.scala
+++ b/tests/src/test/scala/flatgraph/traversal/TraversalTests.scala
@@ -8,7 +8,7 @@ import org.scalatest.wordspec.AnyWordSpec
 import testdomains.generic.GenericDomain
 import testdomains.generic.language.*
 import testdomains.generic.edges.ConnectedTo
-import testdomains.generic.nodes.NodeA
+import testdomains.generic.nodes.{NewNodeA, NewNodeB, NodeA, NodeB}
 
 import scala.collection.mutable
 
@@ -94,7 +94,16 @@ class TraversalTests extends AnyWordSpec {
     oneToFour.without(Set(2, 4)).l shouldBe Seq(1, 3)
   }
 
-  ".help step" should {
+  "collectAll step should collect (and cast) all elements of the given type" in {
+    val anyIter: Iterator[Any] = Iterator("a", "b", 1, 2, 3)
+    val stringIter             = anyIter.collectAll[String]
+
+    // just verifying that we can assign it
+    val stringIter2: Iterator[String] = stringIter
+    stringIter2.size shouldBe 2
+  }
+
+  "help step" should {
     // a specific domain would provide it's own DocSearchPackage implementation, to specify where we're supposed to scan for @Doc annotations
     given DocSearchPackages      = DocSearchPackages.default
     given AvailableWidthProvider = new Table.ConstantWidth(120)
@@ -159,7 +168,6 @@ class TraversalTests extends AnyWordSpec {
       thingTraversalHelpVerbose should include("result to a list")
       thingTraversalHelpVerbose should include("flatgraph.traversal.GenericSt")
     }
-
   }
 
   "generic graph steps" in {