chunk.py

chunk.py

import os
from pathlib import Path
import sys
import numpy as np

def chunk(path, max_lines):

    fd = open(path)
    header = fd.readline()

    dirpath = Path(path.parent, path.name[:-4])

    try:
        os.mkdir(dirpath)
    except:
        pass

    i_chunk = 0
    while True:
        fname_out = f"{path.name[:-4]}_{i_chunk}{path.name[-4:]}"
        fd_out = open(Path(dirpath, fname_out), "w")
        fd_out.write(header)

        stop = False
        for i_line in range(max_lines):
            line = fd.readline()
            if len(line) == 0:
                stop = True
            fd_out.write(line)

        fd_out.close()
        i_chunk += 1

        if stop:
            break

if __name__ == "__main__":
    max_lines = int(sys.argv[1])
    for i in range(2, len(sys.argv)):
        fname = sys.argv[i]
        chunk(Path(fname), max_lines)