fix take_n with more efficient implementation

brilee · Jan 24, 2017 · aaa22dc · aaa22dc
1 parent 6236b20
commit aaa22dc
Showing 1 changed file with 2 additions and 9 deletions.
diff --git a/load_data_sets.py b/load_data_sets.py
@@ -15,15 +15,8 @@
 CHUNK_HEADER_FORMAT = "iii?"
 CHUNK_HEADER_SIZE = struct.calcsize(CHUNK_HEADER_FORMAT)
 
-def take_n(n, iterator):
-    result = []
-    try:
-        for i in range(n):
-            result.append(next(iterator))
-    except StopIteration:
-        pass
-    finally:
-        return result
+def take_n(n, iterable):
+    return list(itertools.islice(iterable, n))
 
 def iter_chunks(chunk_size, iterator):
     while True: