rocksdb/buckifier/util.py

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
"""
This module keeps commonly used components.
"""
from __future__ import absolute_import, division, print_function, unicode_literals

try:
    from builtins import object
except ImportError:
    from __builtin__ import object
import os
import subprocess
import sys
import time


class ColorString(object):
    """Generate colorful strings on terminal"""

    HEADER = "\033[95m"
    BLUE = "\033[94m"
    GREEN = "\033[92m"
    WARNING = "\033[93m"
    FAIL = "\033[91m"
    ENDC = "\033[0m"

    @staticmethod
    def _make_color_str(text, color):
        # In Python2, default encoding for unicode string is ASCII
        if sys.version_info.major <= 2:
            return "".join([color, text.encode("utf-8"), ColorString.ENDC])
        # From Python3, default encoding for unicode string is UTF-8
        return "".join([color, text, ColorString.ENDC])

    @staticmethod
    def ok(text):
        if ColorString.is_disabled:
            return text
        return ColorString._make_color_str(text, ColorString.GREEN)

    @staticmethod
    def info(text):
        if ColorString.is_disabled:
            return text
        return ColorString._make_color_str(text, ColorString.BLUE)

    @staticmethod
    def header(text):
        if ColorString.is_disabled:
            return text
        return ColorString._make_color_str(text, ColorString.HEADER)

    @staticmethod
    def error(text):
        if ColorString.is_disabled:
            return text
        return ColorString._make_color_str(text, ColorString.FAIL)

    @staticmethod
    def warning(text):
        if ColorString.is_disabled:
            return text
        return ColorString._make_color_str(text, ColorString.WARNING)

    is_disabled = False


def run_shell_command(shell_cmd, cmd_dir=None):
    """Run a single shell command.
    @returns a tuple of shell command return code, stdout, stderr"""

    if cmd_dir is not None and not os.path.exists(cmd_dir):
        run_shell_command("mkdir -p %s" % cmd_dir)

    start = time.time()
    print("\t>>> Running: " + shell_cmd)
    p = subprocess.Popen(  # noqa
        shell_cmd,
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        cwd=cmd_dir,
    )
    stdout, stderr = p.communicate()
    end = time.time()

    # Report time if we spent more than 5 minutes executing a command
    execution_time = end - start
    if execution_time > (60 * 5):
        mins = execution_time / 60
        secs = execution_time % 60
        print("\t>time spent: %d minutes %d seconds" % (mins, secs))

    return p.returncode, stdout, stderr


def run_shell_commands(shell_cmds, cmd_dir=None, verbose=False):
    """Execute a sequence of shell commands, which is equivalent to
    running `cmd1 && cmd2 && cmd3`
    @returns boolean indication if all commands succeeds.
    """

    if cmd_dir:
        print("\t=== Set current working directory => %s" % cmd_dir)

    for shell_cmd in shell_cmds:
        ret_code, stdout, stderr = run_shell_command(shell_cmd, cmd_dir)
        if stdout:
            if verbose or ret_code != 0:
                print(ColorString.info("stdout: \n"), stdout)
        if stderr:
            # contents in stderr is not necessarily to be error messages.
            if verbose or ret_code != 0:
                print(ColorString.error("stderr: \n"), stderr)
        if ret_code != 0:
            return False

    return True