samnmax.py 1.82 KB
Newer Older
Maxime Perrotin's avatar
Maxime Perrotin committed
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
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# vim: ai ts=4 sts=4 et sw=4 nu

'''
    This module was developed by Sam & Max
    https://raw.github.com/sametmax/Bat-belt/master/batbelt/hack.py
'''

import sys

from io import BytesIO
from contextlib import contextmanager


@contextmanager
def capture_ouput(stdout_to=None, stderr_to=None):
    """
        Context manager that captures any printed ouput in the 'with' block.

        :Exemple:

        >>> with capture_ouput() as (stdout, stderr):
        ...    print "hello",
        ...
        >>> print stdout.getvalue().upper()
        HELLO
        >>> with capture_ouput() as (stdout, stderr):
            # doctest: +IGNORE_EXCEPTION_DETAIL
        ...    assert False
        ...
        Traceback (most recent call last):
        AssertionError
        >>> from tempfile import NamedTemporaryFile
        >>> f = NamedTemporaryFile(mode="rw+b")
        >>> with capture_ouput(f) as (stdout, stderr):
        ...    print "hello",
        ...
        >>> print stdout.read()
        hello


        .. :warning: this is NOT thread safe.

        .. :note: The file like objects containing the capture are not closed
                  automatically by this context manager. You are responsible
                  to do it.

    It does not capture exception, so they bubble out and print the stack
    trace anyway.
    """

    try:

        stdout, stderr = sys.stdout, sys.stderr
        sys.stdout = c1 = stdout_to or BytesIO()
        sys.stderr = c2 = stderr_to or BytesIO()
        yield c1, c2

    finally:

        sys.stdout = stdout
        sys.stderr = stderr

        try:
            c1.flush()
            c1.seek(0)
        except (ValueError, TypeError):
            pass

        try:
            c2.flush()
            c2.seek(0)
        except (ValueError, TypeError):
            pass