summaryrefslogtreecommitdiffstats
path: root/src/rocksdb/buckifier/util.py
blob: 8943fed2bdfda9b687c70bedab70f90d23af3ee5 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
# 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