1# -*- coding: UTF-8 -*-
2"""
3Capture output (stdout, stderr), logs, etc.
4"""
5
6from __future__ import absolute_import
7from contextlib import contextmanager
8import sys
9from six import StringIO, PY2
10from behave.log_capture import LoggingCapture
11from behave.textutil import text as _text
12
13def add_text_to(value, more_text, separator="\n"):
14    if more_text:
15        if value:
16            if separator and not value.endswith(separator):
17                value += separator
18            value += more_text
19        else:
20            value = more_text
21    return value
22
23
24class Captured(object):
25    """Stores and aggregates captured output data."""
26    empty = u""
27    linesep = u"\n"
28
29    def __init__(self, stdout=None, stderr=None, log_output=None):
30        self.stdout = stdout or self.empty
31        self.stderr = stderr or self.empty
32        self.log_output = log_output or self.empty
33
34    def reset(self):
35        self.stdout = self.empty
36        self.stderr = self.empty
37        self.log_output = self.empty
38
39    # -- PYTHON2:
40    if PY2:
41        def __nonzero__(self):
42            return bool(self.stdout or self.stderr or self.log_output)
43    else:
44        def __bool__(self):
45            return bool(self.stdout or self.stderr or self.log_output)
46
47    @property
48    def output(self):
49        """Makes a simple report of the captured data by concatenating
50        all parts.
51        """
52        output_text = self.stdout
53        output_text = add_text_to(output_text, self.stderr)
54        output_text = add_text_to(output_text, self.log_output)
55        return output_text
56
57    def add(self, captured):
58        """Adds/appends captured output data to this object.
59
60        :param captured:    Captured object whose data should be added.
61        :return: self, to allow daisy-chaining (if needed).
62        """
63        assert isinstance(captured, Captured)
64        self.stdout = add_text_to(self.stdout, captured.stdout, self.linesep)
65        self.stderr = add_text_to(self.stderr, captured.stderr, self.linesep)
66        self.log_output = add_text_to(self.log_output, captured.log_output,
67                                      self.linesep)
68        return self
69
70    def make_report(self):
71        """Makes a detailled report of the captured output data.
72
73        :returns: Report as string.
74        """
75        report_parts = []
76        if self.stdout:
77            parts = ["Captured stdout:", _text(self.stdout).rstrip(), ""]
78            report_parts.extend(parts)
79        if self.stderr:
80            parts = ["Captured stderr:", _text(self.stderr).rstrip(), ""]
81            report_parts.extend(parts)
82        if self.log_output:
83            parts = ["Captured logging:", _text(self.log_output)]
84            report_parts.extend(parts)
85        return self.linesep.join(report_parts).strip()
86
87    def __add__(self, other):
88        """Supports incremental add::
89
90            captured1 = Captured("Hello")
91            captured2 = Captured("World")
92            captured3 = captured1 + captured2
93            assert captured3.stdout == "Hello\nWorld"
94        """
95        new_data = Captured(self.stdout, self.stderr, self.log_output)
96        return new_data.add(other)
97
98    def __iadd__(self, other):
99        """Supports incremental add::
100
101            captured1 = Captured("Hello")
102            captured2 = Captured("World")
103            captured1 += captured2
104            assert captured1.stdout == "Hello\nWorld"
105        """
106        return self.add(other)
107
108
109class CaptureController(object):
110    """Simplifies the lifecycle to capture output from various sources."""
111    def __init__(self, config):
112        self.config = config
113        self.stdout_capture = None
114        self.stderr_capture = None
115        self.log_capture = None
116        self.old_stdout = None
117        self.old_stderr = None
118
119    @property
120    def captured(self):
121        """Provides access of the captured output data.
122
123        :return: Object that stores the captured output parts (as Captured).
124        """
125        stdout = None
126        stderr = None
127        log_out = None
128        if self.config.stdout_capture and self.stdout_capture:
129            stdout = _text(self.stdout_capture.getvalue())
130        if self.config.stderr_capture and self.stderr_capture:
131            stderr = _text(self.stderr_capture.getvalue())
132        if self.config.log_capture and self.log_capture:
133            log_out = _text(self.log_capture.getvalue())
134        return Captured(stdout, stderr, log_out)
135
136    def setup_capture(self, context):
137        assert context is not None
138        if self.config.stdout_capture:
139            self.stdout_capture = StringIO()
140            context.stdout_capture = self.stdout_capture
141
142        if self.config.stderr_capture:
143            self.stderr_capture = StringIO()
144            context.stderr_capture = self.stderr_capture
145
146        if self.config.log_capture:
147            self.log_capture = LoggingCapture(self.config)
148            self.log_capture.inveigle()
149            context.log_capture = self.log_capture
150
151    def start_capture(self):
152        if self.config.stdout_capture:
153            # -- REPLACE ONLY: In non-capturing mode.
154            if not self.old_stdout:
155                self.old_stdout = sys.stdout
156                sys.stdout = self.stdout_capture
157            assert sys.stdout is self.stdout_capture
158
159        if self.config.stderr_capture:
160            # -- REPLACE ONLY: In non-capturing mode.
161            if not self.old_stderr:
162                self.old_stderr = sys.stderr
163                sys.stderr = self.stderr_capture
164            assert sys.stderr is self.stderr_capture
165
166    def stop_capture(self):
167        if self.config.stdout_capture:
168            # -- RESTORE ONLY: In capturing mode.
169            if self.old_stdout:
170                sys.stdout = self.old_stdout
171                self.old_stdout = None
172            assert sys.stdout is not self.stdout_capture
173
174        if self.config.stderr_capture:
175            # -- RESTORE ONLY: In capturing mode.
176            if self.old_stderr:
177                sys.stderr = self.old_stderr
178                self.old_stderr = None
179            assert sys.stderr is not self.stderr_capture
180
181    def teardown_capture(self):
182        if self.config.log_capture:
183            self.log_capture.abandon()
184
185    def make_capture_report(self):
186        """Combine collected output and return as string."""
187        return self.captured.make_report()
188        # report = u""
189        # if self.config.stdout_capture and self.stdout_capture:
190        #     output = self.stdout_capture.getvalue()
191        #     if output:
192        #         output = _text(output)
193        #         report += u"\nCaptured stdout:\n" + output
194        # if self.config.stderr_capture and self.stderr_capture:
195        #     output = self.stderr_capture.getvalue()
196        #     if output:
197        #         output = _text(output)
198        #         report += u"\nCaptured stderr:\n" + output
199        # if self.config.log_capture and self.log_capture:
200        #     output = self.log_capture.getvalue()
201        #     if output:
202        #         output = _text(output)
203        #         report += u"\nCaptured logging:\n" + output
204        # return report
205
206# -----------------------------------------------------------------------------
207# UTILITY FUNCTIONS:
208# -----------------------------------------------------------------------------
209@contextmanager
210def capture_output(controller, enabled=True):
211    """Provides a context manager that starts capturing output
212
213    .. code-block::
214
215        with capture_output(capture_controller):
216            ... # Do something
217    """
218    if enabled:
219        try:
220            controller.start_capture()
221            yield
222        finally:
223            controller.stop_capture()
224    else:
225        # -- CAPTURING OUTPUT is disabled.
226        # Needed to prevent recursive captures with context.execute_steps()
227        yield
228