promptflow

1
import logging
2
import os
3
import subprocess
4
import sys
5
import time
6
import traceback
7

8
module_logger = logging.getLogger(__name__)
9

10

11
class Color:
12
    PURPLE = "\033[95m"
13
    CYAN = "\033[96m"
14
    DARKCYAN = "\033[36m"
15
    BLUE = "\033[94m"
16
    GREEN = "\033[92m"
17
    YELLOW = "\033[93m"
18
    RED = "\033[91m"
19
    BOLD = "\033[1m"
20
    UNDERLINE = "\033[4m"
21
    END = "\033[0m"
22

23

24
def print_red(message):
25
    print(Color.RED + message + Color.END)
26

27

28
def print_blue(message):
29
    print(Color.BLUE + message + Color.END)
30

31

32
def get_test_files(testpath):
33
    if os.path.isfile(testpath):
34
        return [testpath]
35
    else:
36
        res = []
37
        for root, dirs, files in os.walk(testpath):
38
            module_logger.debug("Searching %s for files ending in 'tests.py'", root)
39
            res.extend([os.path.join(root, file) for file in files if file.endswith("tests.py")])
40
        return res
41

42

43
def retry(fn, num_attempts=3):
44
    if num_attempts <= 0:
45
        raise Exception("Illegal num_attempts: {}".format(num_attempts))
46
    count = 0
47
    for _ in range(0, num_attempts):
48
        try:
49
            return fn()
50
        except Exception:
51
            count += 1
52
            print("Execution failed on attempt {} out of {}".format(count, num_attempts))
53
            print("Exception trace:")
54
            traceback.print_exc()
55
            if count == num_attempts:
56
                print("Execution failed after {} attempts".format(count))
57
                raise
58

59

60
def _run_command(
61
    commands,
62
    cwd=None,
63
    stderr=subprocess.STDOUT,
64
    shell=False,
65
    env=None,
66
    stream_stdout=True,
67
    throw_on_retcode=True,
68
    logger=None,
69
):
70
    if logger is None:
71
        logger = module_logger
72

73
    if cwd is None:
74
        cwd = os.getcwd()
75

76
    t0 = time.perf_counter()
77
    try:
78
        logger.debug("[RunCommand]Executing {0} in {1}".format(commands, cwd))
79
        out = ""
80
        p = subprocess.Popen(commands, stdout=subprocess.PIPE, stderr=stderr, cwd=cwd, shell=shell, env=env)
81
        for line in p.stdout:
82
            line = line.decode("utf-8").rstrip()
83
            if line and line.strip():
84
                logger.debug(line)
85
                if stream_stdout:
86
                    sys.stdout.write(line)
87
                    sys.stdout.write("\n")
88
                out += line
89
                out += "\n"
90
        p.communicate()
91
        retcode = p.poll()
92
        if throw_on_retcode:
93
            if retcode:
94
                raise subprocess.CalledProcessError(retcode, p.args, output=out, stderr=p.stderr)
95
        return retcode, out
96
    finally:
97
        t1 = time.perf_counter()
98
        logger.debug("[RunCommand] Execution took {0}s for {1} in {2}".format(t1 - t0, commands, cwd))
99

100

101
def run_command(
102
    commands, cwd=None, stderr=subprocess.STDOUT, shell=False, stream_stdout=True, throw_on_retcode=True, logger=None
103
):
104
    return _run_command(
105
        commands,
106
        cwd=cwd,
107
        stderr=stderr,
108
        shell=shell,
109
        stream_stdout=stream_stdout,
110
        throw_on_retcode=throw_on_retcode,
111
        logger=logger,
112
    )
113