1import collections
2import contextlib
3import errno
4import io
5import os
6import shutil
7import subprocess
8import sys
9import tempfile
10
11from pyperf._utils import popen_communicate, popen_killer
12
13
14@contextlib.contextmanager
15def _capture_stream(name):
16    old_stream = getattr(sys, name)
17    try:
18        if sys.version_info >= (3,):
19            stream = io.StringIO()
20        else:
21            stream = io.BytesIO()
22        setattr(sys, name, stream)
23        yield stream
24    finally:
25        setattr(sys, name, old_stream)
26
27
28def capture_stdout():
29    return _capture_stream('stdout')
30
31
32def capture_stderr():
33    return _capture_stream('stderr')
34
35
36@contextlib.contextmanager
37def temporary_file(**kwargs):
38    tmp_filename = tempfile.mktemp(**kwargs)
39    try:
40        yield tmp_filename
41    finally:
42        try:
43            os.unlink(tmp_filename)
44        except OSError as exc:
45            if exc.errno != errno.ENOENT:
46                raise
47
48
49@contextlib.contextmanager
50def temporary_directory():
51    tmpdir = tempfile.mkdtemp()
52    try:
53        yield tmpdir
54    finally:
55        shutil.rmtree(tmpdir)
56
57
58def benchmark_as_json(benchmark, compact=True):
59    with temporary_file() as tmp_name:
60        benchmark.dump(tmp_name, compact=compact)
61        with io.open(tmp_name, 'r', encoding='utf-8') as tmp:
62            return tmp.read()
63
64
65def compare_benchmarks(testcase, bench1, bench2):
66    json1 = benchmark_as_json(bench1, compact=False)
67    json2 = benchmark_as_json(bench2, compact=False)
68    testcase.assertEqual(json1, json2)
69
70
71ProcResult = collections.namedtuple('ProcResult', 'returncode stdout stderr')
72
73
74def get_output(cmd, **kw):
75    proc = subprocess.Popen(cmd,
76                            stdout=subprocess.PIPE,
77                            stderr=subprocess.PIPE,
78                            universal_newlines=True,
79                            **kw)
80    stdout, stderr = popen_communicate(proc)
81    return ProcResult(proc.returncode, stdout, stderr)
82
83
84def run_command(cmd, **kw):
85    proc = subprocess.Popen(cmd, **kw)
86    with popen_killer(proc):
87        proc.wait()
88    return proc.returncode
89