util.py

from datetime import timedelta
import functools
import logging
from typing import Optional, TypeVar, Union, List, Any

TType = TypeVar("TType")
import os
import sys
import config
import tempfile


def flatten_list(lst):
    return [item for sublist in lst for item in sublist]


def unzip(lst):
    res = [[i for i, j in lst], [j for i, j in lst]]
    return res


def pad(lst, index):
    if index >= len(lst):
        lst += [None] * (index + 1 - len(lst))
    return lst


def print_time_delta(prefix, start_time, end_time):
    ## Always output time in the log.
    time_difference = (end_time - start_time) / timedelta(milliseconds=1)
    ## If output_time flag is set, log the time
    if config.OUTPUT_TIME:
        log("{} time:".format(prefix), time_difference, " ms", level=0)
    else:
        log("{} time:".format(prefix), time_difference, " ms")


## This function decorates a function to add the logging prefix (without missing the old prefix)
def logging_prefix(logging_prefix):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            old_prefix = config.LOGGING_PREFIX
            config.LOGGING_PREFIX = logging_prefix
            result = func(*args, **kwargs)
            config.LOGGING_PREFIX = old_prefix
            return result

        return wrapper

    return decorator


## This is a wrapper for prints
def log(*args, end="\n", level=1):
    ## If the debug logging level is at least
    ## as high as this log message.
    ## TODO: Allow all levels
    if level >= 1:
        concatted_args = " ".join([str(a) for a in list(args)])
        logging.info(f"{config.LOGGING_PREFIX} {concatted_args}")


def ptempfile():
    fd, name = tempfile.mkstemp(dir=config.PASH_TMP_PREFIX)
    ## TODO: Get a name without opening the fd too if possible
    os.close(fd)
    return name


def return_empty_list_if_none_else_itself(
    arg: Optional[TType],
) -> Union[TType, List[Any]]:  # list always empty
    if arg is None:
        return []
    else:
        return arg


def return_default_if_none_else_itself(arg: Optional[TType], default: TType) -> TType:
    if arg is None:
        return default
    else:
        return arg


## This function gets a key and a value from the ast json format
def get_kv(dic):
    return (dic[0], dic[1])


def make_kv(key, val):
    return [key, val]