utils.py

import argparse
import datetime
import json
import os


def conditional_value(var, nul, default):
    """
    Set the value to default if nul.
    Parameters
    ----------
    var : any
        The value
    nul : any
        The null value. Assigns default if var == nul
    default : any
        The default value
    Returns
    -------
    type(var)
        The final value
    """
    if var != nul:
        return var
    else:
        return default


def remove_keys(d, keys_to_remove):
    """
    Removes given keys from the dict. Returns a new list.
    Parameters
    ----------
    d : dict
        The initial dictionary
    keys_to_remove : list
        List of keys to remove from dict
    Returns
    -------
    dict
        A new dictionary with the given keys removed.
    """
    return {key: d[key] for key in d if key not in keys_to_remove}


def get_args():
    """
    Utility to parse arguments.
    Returns
    -------
    args
        Command line arguments
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("-mid", "--machine_id", type=int, default=0)
    parser.add_argument("-ps", "--procs_per_machine", type=int, default=1)
    parser.add_argument("-ms", "--machines", type=int, default=1)
    parser.add_argument(
        "-ld",
        "--log_dir",
        type=str,
        default="./{}".format(datetime.datetime.now().isoformat(timespec="minutes")),
    )
    parser.add_argument("-is", "--iterations", type=int, default=1)
    parser.add_argument("-cf", "--config_file", type=str, default="config.ini")
    parser.add_argument("-ll", "--log_level", type=str, default="INFO")
    parser.add_argument("-gf", "--graph_file", type=str, default="36_nodes.edges")
    parser.add_argument("-gt", "--graph_type", type=str, default="edges")
    parser.add_argument("-ta", "--test_after", type=int, default=5)

    args = parser.parse_args()
    return args


def write_args(args, path):
    """
    Write arguments to a json file
    Parameters
    ----------
    args : args
        Command line args
    path : str
        Location of the file to write to
    """
    data = {
        "machine_id": args.machine_id,
        "procs_per_machine": args.procs_per_machine,
        "machines": args.machines,
        "log_dir": args.log_dir,
        "iterations": args.iterations,
        "config_file": args.config_file,
        "log_level": args.log_level,
        "graph_file": args.graph_file,
        "graph_type": args.graph_type,
        "test_after": args.test_after,
    }
    with open(os.path.join(path, "args.json"), "w") as of:
        json.dump(data, of)