python_util/shell_util.py at master · didicout/python_util · GitHub

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
# -*- coding: utf-8 -*-
"""
shell util
"""
from __future__ import print_function
import subprocess
import sys


class Color:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


def run_bash_script(script,
                    print_trace=False,
                    return_stdout=False,
                    return_stderr=False,
                    combine_stderr=False):
    """
    run bash script
    :param script: bash script string
    :param print_trace: whether print script and result
    :param return_stdout: catch stdout as string and return instead of printing to sys.stdout
    :param return_stderr: catch stderr as string and return instead of printing to sys.stderr
    :param combine_stderr: redirect stderr to stdout
    :type script: str
    :type print_trace: bool
    :type return_stdout: bool
    :type return_stderr: bool
    :type combine_stderr: bool
    :return: ret_value, stdout, stderr
    :rtype: tuple[int, str, str]
    """
    if print_trace:
        print(Color.WARNING + script + Color.ENDC)
    if return_stdout:
        target_stdout = subprocess.PIPE
    else:
        target_stdout = sys.stdout
    if combine_stderr:
        target_stderr = subprocess.STDOUT
    else:
        if return_stderr:
            target_stderr = subprocess.PIPE
        else:
            target_stderr = sys.stderr

    p = subprocess.Popen("/bin/bash", stdout=target_stdout, stderr=target_stderr,
                         stdin=subprocess.PIPE)
    p.stdin.write(script)
    ret_stdout = None
    ret_stderr = None
    tmp_stdout, tmp_stderr = p.communicate()
    ret_value = p.returncode
    if return_stdout:
        ret_stdout = tmp_stdout
    if return_stderr and not combine_stderr:
        ret_stderr = tmp_stderr
    if print_trace:
        if ret_value == 0:
            print(Color.OKGREEN + str(ret_value) + Color.ENDC)
        else:
            print(Color.FAIL + str(ret_value) + Color.ENDC)
        if return_stdout:
            print(ret_stdout)
        if return_stderr and not combine_stderr:
            print(ret_stderr)
    return ret_value, ret_stdout, ret_stderr


# test
print("================= 1 =================")
print(run_bash_script("echo 123"))
print("================= 2 =================")
print(run_bash_script("echo 123 && cp", return_stderr=True, return_stdout=True, print_trace=True))
print("================= 3 =================")
print(run_bash_script("echo 123 && cp", return_stdout=True, return_stderr=True, combine_stderr=True, print_trace=True))
print("================= 4 =================")
print(run_bash_script("echo 123 && cp", return_stdout=True, return_stderr=False, combine_stderr=True, print_trace=True))
print("================= 5 =================")
print(run_bash_script("echo 123 && cp", return_stdout=True, return_stderr=False, print_trace=True))
print("================= 6 =================")
print(run_bash_script("echo 123 && cp", return_stdout=False, return_stderr=True, combine_stderr=True, print_trace=True))
print("================= 7 =================")
print(run_bash_script("echo 123 && cp", return_stdout=False, return_stderr=True, combine_stderr=False, print_trace=True))
print("================= 8 =================")
print(run_bash_script("echo 123 && cp", return_stdout=False, return_stderr=False, combine_stderr=True, print_trace=True))
print("================= 9 =================")
print(run_bash_script("echo 123 && cp", return_stdout=False, return_stderr=False, combine_stderr=False, print_trace=True))