eigenangles/scripts/merge_dataframes.py

import argparse
import pandas as pd
import numpy as np
from pathlib import Path
from utils import load_df

def use_file(file, include=[], exclude=[]):
    is_df = Path(file).suffix == '.csv'
    excluded = np.array([excl in file.name for excl in exclude]).any()
    included = np.array([incl in file.name for incl in include]).all()
    return is_df and included and not excluded

if __name__ == '__main__':
    CLI = argparse.ArgumentParser()
    CLI.add_argument("--input", nargs='?', type=Path, default=None)
    CLI.add_argument("--output", nargs='?', type=Path)
    CLI.add_argument("--exclude", nargs='?', type=lambda s: s.split(' '),
                     default=['rewiring_results'])
    args, unknown = CLI.parse_known_args()

    dfs = []
    if args.input is None:
        args.input = args.output.parent

    if args.input.is_file():
        with open(args.input, "r") as file:
            for path in file.readlines():
                dfs.append(load_df(path.strip('\n')))
    elif args.input.is_dir():
        for file in args.input.rglob("*"):
            if use_file(file, exclude=args.exclude):
                dfs.append(load_df(file))

    full_df = pd.concat(dfs, ignore_index=True)
    full_df.to_csv(args.output)