Source code for research.utils._data

"""
Data I/O utils. Later on I might add other data handling utilities.
"""
from os import listdir
from os.path import isdir, join
import pandas as pd
from sqlite3 import connect


def load_datasets(data_dir, suffix="", target_exists=True, **read_csv_kwargs):
    """Load datasets from sqlite database and/or csv files."""
    assert isdir(data_dir), "`data_dir` must be a directory."

    # Filter data by suffix
    dat_names = [dat for dat in listdir(data_dir) if dat.endswith(suffix)]

    # Read data
    datasets = []
    for dat_name in dat_names:
        data_path = join(data_dir, dat_name)

        # Handle csv data
        if dat_name.endswith(".csv"):
            ds = pd.read_csv(data_path, **read_csv_kwargs)
            name = dat_name.replace(".csv", "").replace("_", " ").upper()
            if target_exists:
                ds = (ds.iloc[:, :-1], ds.iloc[:, -1])
            datasets.append((name, ds))

        # Handle sqlite database
        elif dat_name.endswith(".db"):
            with connect(data_path) as connection:
                datasets_names = [
                    name[0]
                    for name in connection.execute(
                        "SELECT name FROM sqlite_master WHERE type='table';"
                    )
                ]
                for dataset_name in datasets_names:
                    ds = pd.read_sql(f'select * from "{dataset_name}"', connection)
                    if target_exists:
                        ds = (ds.iloc[:, :-1], ds.iloc[:, -1])
                    datasets.append((dataset_name.replace("_", " ").upper(), ds))
    return datasets