from numbers import Integral

import pandas as pd


def get_colnames(cursor):
    """ Get column names for cursor's result set """
    return [coldesc[0] for coldesc in cursor.description]


# TODO: have the batched dict and pd methods emit batches instead of flattening it out (otherwise
# what's the point of the batching....)

def read_cursor(cursor, batch_size='all'):
    """ Consistent interface to a DBAPI-compliant Cursor object """
    if batch_size == 'all':
        yield from cursor.fetchall()
    elif batch_size == 1:
        while True:
            row = cursor.fetchone()
            if row is not None:
                yield row
            else:
                break
    elif isinstance(batch_size, Integral) and batch_size > 0:
        while True:
            rows = cursor.fetchmany(batch_size)
            if len(rows):
                yield from rows
            else:
                break
    else:
        raise ValueError('Invalid batch size')


def read_cursor_dicts(cursor, batch_size='all'):
    """ Read cursor, with each row as a dictionary """
    colnames = get_colnames(cursor)
    for row in read_cursor(cursor, batch_size):
        yield dict(zip(colnames, row))


def read_cursor_pandas(cursor, batch_size='all'):
    colnames = get_colnames(cursor)
    return pd.DataFrame(list(read_cursor(cursor, batch_size)), columns=colnames)