From b5f1e2a3e7c9189afa1998389ca0262c8e3b775d Mon Sep 17 00:00:00 2001
From: Jeremy Fincher <jemfinch@users.sourceforge.net>
Date: Mon, 30 May 2005 19:19:11 +0000
Subject: [PATCH] Added utils.iter.limited, an iterable that limits the number
 of elements that can be taken from another iterable.

---
 src/utils/iter.py  | 11 +++++++++++
 test/test_utils.py |  8 ++++++++
 2 files changed, 19 insertions(+)

diff --git a/src/utils/iter.py b/src/utils/iter.py
index 3102a9a37..750df99ca 100644
--- a/src/utils/iter.py
+++ b/src/utils/iter.py
@@ -158,4 +158,15 @@ def startswith(long, short):
     except StopIteration:
         return True
 
+def limited(iterable, limit):
+    i = limit
+    iterable = iter(iterable)
+    try:
+        while i:
+            yield iterable.next()
+            i -= 1
+    except StopIteration:
+        raise ValueError, 'Expected %s elements in iterable (%r), got %s.' % \
+              (limit, iterable, limit-i)
+
 # vim:set shiftwidth=4 tabstop=4 expandtab textwidth=79:
diff --git a/test/test_utils.py b/test/test_utils.py
index 3c4e0835d..731396aa0 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -384,6 +384,14 @@ class StrTest(SupyTestCase):
 
 
 class IterTest(SupyTestCase):
+    def testLimited(self):
+        L = range(10)
+        self.assertEqual([], list(utils.iter.limited(L, 0)))
+        self.assertEqual([0], list(utils.iter.limited(L, 1)))
+        self.assertEqual([0, 1], list(utils.iter.limited(L, 2)))
+        self.assertEqual(range(10), list(utils.iter.limited(L, 10)))
+        self.assertRaises(ValueError, list, utils.iter.limited(L, 11))
+        
     def testRandomChoice(self):
         choice = utils.iter.choice
         self.assertRaises(IndexError, choice, {})