Refactored to use itertools

Found an itertools mechanism that works for iterating
through the current and next row in the cursor.
This commit is contained in:
Simon Willison 2019-01-26 10:58:45 -08:00
commit 3f2e711a4b
2 changed files with 5 additions and 32 deletions

View file

@ -1,6 +1,6 @@
import click
import sqlite_utils
from sqlite_utils.utils import iter_pairs
import itertools
import json as json_std
import sys
import csv as csv_std
@ -140,11 +140,12 @@ def json(path, sql, nl, arrays):
# We have to iterate two-at-a-time so we can know if we
# should output a trailing comma or if we have reached
# the last row.
row = None
current_iter, next_iter = itertools.tee(cursor, 2)
next(next_iter, None)
first = True
headers = [c[0] for c in cursor.description]
for row, next_row, is_last in iter_pairs(cursor):
# We now reliably have row and next_row
for row, next_row in itertools.zip_longest(current_iter, next_iter):
is_last = next_row is None
data = row
if not arrays:
data = dict(zip(headers, row))

View file

@ -1,28 +0,0 @@
def iter_pairs(iterator):
# Yields (item, next_item, is_last) pairs
# Last row is (item, None, True)
first = True
next_item = None
while True:
next_done = False
if first:
item, done = next_with_done(iterator)
if done:
break
next_item, next_done = next_with_done(iterator)
first = False
else:
item = next_item
next_item, next_done = next_with_done(iterator)
yield item, next_item, next_done
if next_done:
break
def next_with_done(iterator):
try:
return next(iterator), False
except StopIteration:
return None, True