diff options
-rw-r--r-- | core/src/main/scala/spark/util/CleanupIterator.scala | 25 |
1 files changed, 25 insertions, 0 deletions
diff --git a/core/src/main/scala/spark/util/CleanupIterator.scala b/core/src/main/scala/spark/util/CleanupIterator.scala new file mode 100644 index 0000000000..d2093c0230 --- /dev/null +++ b/core/src/main/scala/spark/util/CleanupIterator.scala @@ -0,0 +1,25 @@ +package spark.util + +/** + * Wrapper around an iterator which calls a cleanup method when its finished iterating through its elements + */ +abstract class CleanupIterator[+A, +I <: Iterator[A]](sub: I) extends Iterator[A]{ + def next = sub.next + def hasNext = { + val r = sub.hasNext + if (!r) { + cleanup + } + r + } + + def cleanup +} + +object CleanupIterator { + def apply[A, I <: Iterator[A]](sub: I, cleanupFunction: => Unit) : CleanupIterator[A,I] = { + new CleanupIterator[A,I](sub) { + def cleanup = cleanupFunction + } + } +}
\ No newline at end of file |