loneElement and loneElementOption: extension methods for Iterable

Author: mpollmeier

codepropertygraph/src/main/scala/io/shiftleft/Implicits.scala

@@ -3,45 +3,85 @@ package io.shiftleft
 import org.slf4j.{Logger, LoggerFactory}
 
 object Implicits {
-
   private val logger: Logger = LoggerFactory.getLogger(Implicits.getClass)
 
-  implicit class IterableOnceDeco[T](val iterable: IterableOnce[T]) extends AnyVal {
-    def onlyChecked: T = {
-      if (iterable.iterator.hasNext) {
-        val res = iterable.iterator.next()
-        if (iterable.iterator.hasNext) {
-          logger.warn("iterator was expected to have exactly one element, but it actually has more")
+  extension [A](iterable: IterableOnce[A]) {
+
+    /** @see {{{loneElement(hint)}}} */
+    def loneElement: A =
+      loneElement(hint = "")
+
+    /** @return
+      *   the one and only element from an Iterable
+      * @throws NoSuchElementException
+      *   if the Iterable is empty
+      * @throws AssertionError
+      *   if the Iterable has more than one element
+      */
+    def loneElement(hint: String): A = {
+      lazy val hintMaybe =
+        if (hint.isEmpty) ""
+        else s" Hint: $hint"
+
+      val iter = iterable.iterator
+      if (iter.isEmpty) {
+        throw new NoSuchElementException(
+          s"Iterable was expected to have exactly one element, but it is empty.$hintMaybe"
+        )
+      } else {
+        val res = iter.next()
+        if (iter.hasNext) {
+          val collectionSizeHint = iterable.knownSize match {
+            case -1 => "it has more than one" // cannot be computed cheaply, i.e. without traversing the collection
+            case knownSize => s"it has $knownSize"
+          }
+          throw new AssertionError(
+            s"Iterable was expected to have exactly one element, but $collectionSizeHint.$hintMaybe"
+          )
         }
         res
-      } else { throw new NoSuchElementException() }
+      }
     }
-  }
 
-  /** A wrapper around a Java iterator that throws a proper NoSuchElementException.
-    *
-    * Proper in this case means an exception with a stack trace. This is intended to be used as a replacement for next()
-    * on the iterators returned from TinkerPop since those are missing stack traces.
-    */
-  implicit class JavaIteratorDeco[T](val iterator: java.util.Iterator[T]) extends AnyVal {
-    def nextChecked: T = {
-      try {
-        iterator.next
-      } catch {
-        case _: NoSuchElementException =>
-          throw new NoSuchElementException()
+    /** @see {{{loneElementOption(hint)}}} */
+    def loneElementOption: Option[A] =
+      loneElementOption(hint = None)
+
+    /** @return
+      *   {{{Some(element)}}} if the Iterable has exactly one element, or {{{None}}} if the Iterable has zero or more
+      *   than 1 element. Note: if the lone element is {{{null}}}, this will return {{{Some(null)}}}, which is in
+      *   accordance with how {{{headOption}}} works.
+      */
+    def loneElementOption(hint: String | None.type = None): Option[A] = {
+      val iter = iterable.iterator
+      if (iter.isEmpty) {
+        None
+      } else {
+        val result = iter.next()
+        if (iter.hasNext) {
+          None
+        } else {
+          Some(result)
+        }
       }
     }
+  }
 
+  implicit class IterableOnceDeco[T](val iterable: IterableOnce[T]) extends AnyVal {
+    @deprecated(
+      "please use `loneElement` instead, which has a better name and will throw if the iterable has more than one element (rather than just log.warn)",
+      since = "1.7.42 (July 2025)"
+    )
     def onlyChecked: T = {
-      if (iterator.hasNext) {
-        val res = iterator.next
-        if (iterator.hasNext) {
+      if (iterable.iterator.hasNext) {
+        val res = iterable.iterator.next()
+        if (iterable.iterator.hasNext) {
           logger.warn("iterator was expected to have exactly one element, but it actually has more")
         }
         res
-      } else { throw new NoSuchElementException() }
+      } else {
+        throw new NoSuchElementException()
+      }
     }
   }
-
 }

codepropertygraph/src/test/scala/io/shiftleft/ImplicitsTests.scala

@@ -0,0 +1,52 @@
+package io.shiftleft
+
+import io.shiftleft.Implicits.*
+import org.scalatest.matchers.should.Matchers
+import org.scalatest.wordspec.AnyWordSpec
+
+import scala.collection.mutable.ArrayBuffer
+
+class ImplicitsTests extends AnyWordSpec with Matchers {
+
+  "loneElement returns the one and only element from an Iterable, and throws an exception otherwise" in {
+    Seq(1).loneElement shouldBe 1
+    Seq(1).loneElement("some context") shouldBe 1
+    Seq(null).loneElement shouldBe null
+
+    intercept[NoSuchElementException] {
+      Seq.empty.loneElement
+    }.getMessage should include("it is empty")
+
+    intercept[NoSuchElementException] {
+      Seq.empty.loneElement("some context")
+    }.getMessage should include("it is empty. Hint: some context")
+
+    intercept[AssertionError] {
+      Seq(1, 2).loneElement
+    }.getMessage should include("it has more than one")
+
+    intercept[AssertionError] {
+      ArrayBuffer(1, 2).loneElement
+    }.getMessage should include(
+      "it has 2"
+    ) // ArrayBuffer can 'cheaply' compute their size, so we can have it in the exception message
+
+    intercept[AssertionError] {
+      Seq(1, 2).loneElement("some context")
+    }.getMessage should include("it has more than one. Hint: some context")
+  }
+
+  "loneElementOption returns an Option of the one and only element from an Iterable, or else None" in {
+    Seq(1).loneElementOption shouldBe Some(1)
+    Seq(1).loneElementOption("some context") shouldBe Some(1)
+    Seq(null).loneElementOption shouldBe Some(null)
+    Seq(null).loneElementOption("some context") shouldBe Some(null)
+
+    Seq.empty.loneElementOption shouldBe None
+    Seq.empty.loneElementOption("some context") shouldBe None
+
+    Seq(1, 2).loneElementOption shouldBe None
+    Seq(1, 2).loneElementOption("some context") shouldBe None
+  }
+
+}

project/build.properties

@@ -1 +1 @@
-sbt.version=1.10.11
+sbt.version=1.11.2