utils.py

import os
import sys
import subprocess
import json
import pickle
import xml.etree.ElementTree
import logging
import pandas as pd
import difflib
from functools import wraps
import time
import hashlib
import concurrent.futures
import signal
import multiprocessing


def get_logger_for_file(loglevel=logging.ERROR, logname=__name__):
    logger = logging.getLogger(logname)
    if not logger.hasHandlers():
        __formatter = logging.Formatter('%(asctime)s %(levelname)s %(funcName)s : %(message)s')
        __handler = logging.StreamHandler(sys.stderr)
        __handler.setFormatter(__formatter)
        logger.addHandler(__handler)
    logger.setLevel(loglevel)
    return logger


loglevel = logging.DEBUG
_logger = get_logger_for_file(loglevel)
# time logger is always in debug mode
_time_logger = get_logger_for_file(logging.DEBUG, 'timing')


def run_cmd(cmd, shell=True, timeout=0.0):
    _logger.info(f'run cmd: {cmd}')
    # ret, output = subprocess.getstatusoutput(cmd, encoding='utf-8')
    try:
        ret = 0
        if timeout <= 0.0:
            output = subprocess.check_output(cmd,
                                             shell=shell,
                                             stderr=subprocess.STDOUT)
        else:
            output = subprocess.check_output(cmd,
                                             shell=shell,
                                             stderr=subprocess.STDOUT,
                                             timeout=timeout)
    except subprocess.CalledProcessError as e:
        _logger.error(f'Error: {cmd}')
        ret = e.returncode
        output = e.output
    except subprocess.TimeoutExpired as e:
        _logger.error(f'Timeout: {cmd}')
        ret = -1
        output = e.output
    return ret, output.decode('utf-8')


def init_dir_for_dump(dump_path):
    if '/' not in dump_path and '\\' not in dump_path:
        return dump_path
    dump_dir = os.path.dirname(dump_path)
    if not os.path.exists(dump_dir):
        os.makedirs(dump_dir)


def load_json(json_path):
    with open(json_path, 'r') as f:
        data = json.load(f)
    return data


def dump_json(obj, dump_path, indent=1):
    init_dir_for_dump(dump_path)
    with open(dump_path, 'w') as f:
        json.dump(obj, f, indent=indent)


def load_pkl(pickle_path):
    with open(pickle_path, 'rb') as f:
        data = pickle.load(f)
    return data


def dump_pkl(obj, dump_path):
    init_dir_for_dump(dump_path)
    with open(dump_path, 'wb') as f:
        pickle.dump(obj, f, protocol=4)


def load_csv(csv_path, **kwargs):
    try:
        return pd.read_csv(csv_path, **kwargs)
    except Exception as e:
        _logger.error(f'Fail to load {csv_path}')
        raise e


def get_str_similar(str1, str2):
    return difflib.SequenceMatcher(None, str1, str2).ratio()


def timing(f):
    @wraps(f)
    def wrap(*args, **kw):
        ts = time.time()
        result = f(*args, **kw)
        te = time.time()
        cost = te - ts
        if cost < 60:
            _time_logger.debug('func:%r took: %2.4f sec' % \
                (f.__name__, cost))
        elif 60 < cost < 3600:
            _time_logger.debug('func:%r took: %2.4f min' % \
                (f.__name__, cost / 60))
        else:
            _time_logger.debug('func:%r took: %2.4f hours' % \
                (f.__name__, cost / 3600))
        return result
    return wrap


def get_all_loggers():
    loggers = [logging.getLogger(name) for name in logging.root.manager.loggerDict]
    return loggers


def get_file_hash(file_path):
    with open(file_path, 'rb') as f:
        return hashlib.md5(f.read()).hexdigest()


def timeout_run_function(func_ptr, args, timeout):
    """
    Need to catch concurrent.futures.TimeoutError
    """
    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(func_ptr, *args)
        result = future.result(timeout=timeout)
        return result


class TimeoutException(Exception):
    pass


def timeout(seconds=5, error_message="Function call timed out"):
    """
    A decorator that raises a TimeoutException if the decorated function 
    runs longer than `seconds`.
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # Define a signal handler that raises the timeout exception
            def _handle_timeout(signum, frame):
                raise TimeoutException(error_message)
            
            # Set the signal handler and a timeout alarm
            signal.signal(signal.SIGALRM, _handle_timeout)
            signal.alarm(seconds)
            try:
                return func(*args, **kwargs)
            finally:
                # Cancel the alarm if the function finishes in time
                signal.alarm(0)
        return wrapper
    return decorator


def _eval_for_timeout(code, queue):
    try:
        result = eval(code)
        queue.put((True, result))
    except Exception as e:
        queue.put((False, e))


def eval_with_timeout(code, timeout):
    queue = multiprocessing.Queue()
    process = multiprocessing.Process(target=_eval_for_timeout,
                                      args=(code, queue))
    process.start()
    process.join(timeout=timeout)
    if process.is_alive():
        process.terminate()
        process.join()
        raise TimeoutException('Timeout for eval')
    else:
        succ, res = queue.get()
        if succ:
            return res
        else:
            raise res