__all__ = ['partition']

from datetime import datetime
import gzip
from os import makedirs
from pathlib import Path

formats = {
    'day': '%Y/%m/%d',
    'hour': '%Y/%m/%d/%H',
    'minute': '%Y/%m/%d/%H/%M',
}

def partition(lines, directory='.', format=None, granularity='hour', compressed=True):
    current = None
    file = None
    if format is None:
        format = formats[granularity]
    for line in lines:
        now = datetime.utcnow()
        if current is None or getattr(current, granularity) != getattr(now, granularity):
            if file is not None:
                file.close()
            current = now
            path = Path(directory) / f"{current.strftime(format)}.jsonl"
            makedirs(path.parent, exist_ok=True)
            file = gzip.open(str(path) + '.gz', 'wb') if compressed else open(path, 'wb')
        file.write(line)
    if file is not None:
        file.close()
