fix missing

1 files changed, 12 insertions, 9 deletions

diff --git a/split-csv.py b/split-csv.py
index 62dc1597..2db45d85 100644
--- a/split-csv.py
+++ b/split-csv.py
@@ -9,23 +9,30 @@ import click
 
 @click.command()
 @click.option('--count', '-c', default=2, help='Number of subdivisions.')
+@click.option('--has_keys/--has_no_keys', '-k', default=False, help='Whether to split off the keys.')
 @click.option('--shuffle/--no_shuffle', default=False, help='Whether to shuffle.')
 @click.argument('filename')
-def split_csv(count, shuffle, filename):
+def split_csv(count, has_keys, shuffle, filename):
   """Split a CSV into groups."""
   with open(filename, 'r') as f:
     reader = csv.reader(f)
     lines = list(reader)
-    keys = lines[0]
-    lines = lines[1:]
+    if has_keys:
+      keys = lines[0]
+      lines = lines[1:]
+    else:
+      keys = None
 
     fn, ext = os.path.splitext(filename)
     if shuffle:
       random.shuffle(lines)
-    for index, chunk in enumerate(chunks(lines, count)):
+    n = max(1, ceil(len(lines) / count))
+    for index in range(count):
+      m = index * n
+      chunk = lines[m:m+n]
+      print(chunk[0])
       out_fn = fn + '-' + str(index+1) + ext
       write_csv(out_fn, keys, chunk)
-      # sys.exit(1)
 
 # Write a CSV
 def write_csv(fn, keys, chunk):
@@ -37,10 +44,6 @@ def write_csv(fn, keys, chunk):
     for row in chunk:
       writer.writerow(row)
 
-# Split an array into chunks
-def chunks(l, n):
-  n = max(1, ceil(len(l) / n))
-  return (l[i:i+n] for i in range(0, len(l), n))
 
 if __name__ == '__main__':
   split_csv()