sdpython / cpyquickhelper
1
"""
2
@file
3
@brief Python wrapper around C functions.
4
"""
5 5
from typing import Tuple, Union
6 5
from io import StringIO
7 5
from contextlib import redirect_stdout, redirect_stderr
8 5
from .stdchelper import begin_capture, end_capture, get_capture  # pylint: disable=E0611
9

10

11 5
def capture_output_c(function_to_call) -> Tuple[bytes, Union[bytes, None]]:
12
    """
13
    Capture the standard output and error for
14
    function *function_to_call*, it wraps C code which
15
    catches information from the command line.
16

17
    @param      function_to_call        function to call
18
    @return                             output, error
19

20
    This function must not be called in parallel with another
21
    call of the same function.
22

23
    .. warning:: *error* is always empty. Both streams are merged.
24
    """
25
    if not callable(function_to_call):  # pragma no cover
26
        raise TypeError("function_to_call must be callable.")
27 5
    begin_capture()
28 5
    fout = function_to_call()
29 5
    end_capture()
30 5
    res = get_capture()
31 5
    if res is None:
32 0
        return fout, None, None
33
    if isinstance(res, bytes):  # pragma: no cover
34
        return fout, res, None
35
    if isinstance(res, tuple):  # pragma: no cover
36
        return (fout, ) + res
37
    raise TypeError(  # pragma no cover
38
        "Unexpected return type '{0}'.".format(type(res)))
39

40

41 5
def capture_output_py(function_to_call) -> Tuple[str, str]:
42
    """
43
    Capture the standard output and error for
44
    function *function_to_call* with function
45
    `redirect_stdout <https://docs.python.org/3/library/contextlib.html#contextlib.redirect_stdout>`_
46
    and function
47
    `redirect_stderr <https://docs.python.org/3/library/contextlib.html#contextlib.redirect_stderr>`_.
48

49
    @param      function_to_call        function to call
50
    @return                             output, error
51

52
    This function must not be called in parallel with another
53
    call of the same function.
54

55
    .. warning:: *error* is always empty. Both streams are merged.
56
    """
57 5
    if not callable(function_to_call):
58 5
        raise TypeError("function_to_call must be callable.")
59 5
    out, err = StringIO(), StringIO()
60 5
    with redirect_stdout(out):
61 5
        with redirect_stderr(err):
62 5
            fout = function_to_call()
63 5
    return fout, out.getvalue(), err.getvalue()
64

65

66 5
def capture_output(function_to_call, lang="py"):
67
    """
68
    Catch standard output and error for function
69
    *function_to_call*. If lang is *'py'*, calls
70
    @see fn capture_output_py or @see fn capture_output_c
71
    if lang is *'c'*.
72

73
    @param      function_to_call        function to call
74
    @return                             output, error
75
    """
76 5
    if lang == "py":
77 5
        return capture_output_py(function_to_call)
78 5
    elif lang == "c":
79 5
        return capture_output_c(function_to_call)
80 5
    raise ValueError("lang must be 'py' or 'c' not '{0}'".format(lang))

Read our documentation on viewing source code .

Loading