1from _pydevd_bundle.pydevd_constants import ForkSafeLock, get_global_debugger, IS_PY2
2import os
3import sys
4from contextlib import contextmanager
5
6
7class IORedirector:
8    '''
9    This class works to wrap a stream (stdout/stderr) with an additional redirect.
10    '''
11
12    def __init__(self, original, new_redirect, wrap_buffer=False):
13        '''
14        :param stream original:
15            The stream to be wrapped (usually stdout/stderr, but could be None).
16
17        :param stream new_redirect:
18            Usually IOBuf (below).
19
20        :param bool wrap_buffer:
21            Whether to create a buffer attribute (needed to mimick python 3 s
22            tdout/stderr which has a buffer to write binary data).
23        '''
24        self._lock = ForkSafeLock(rlock=True)
25        self._writing = False
26        self._redirect_to = (original, new_redirect)
27        if wrap_buffer and hasattr(original, 'buffer'):
28            self.buffer = IORedirector(original.buffer, new_redirect.buffer, False)
29
30    def write(self, s):
31        # Note that writing to the original stream may fail for some reasons
32        # (such as trying to write something that's not a string or having it closed).
33        with self._lock:
34            if self._writing:
35                return
36            self._writing = True
37            try:
38                for r in self._redirect_to:
39                    if hasattr(r, 'write'):
40                        r.write(s)
41            finally:
42                self._writing = False
43
44    def isatty(self):
45        for r in self._redirect_to:
46            if hasattr(r, 'isatty'):
47                return r.isatty()
48        return False
49
50    def flush(self):
51        for r in self._redirect_to:
52            if hasattr(r, 'flush'):
53                r.flush()
54
55    def __getattr__(self, name):
56        for r in self._redirect_to:
57            if hasattr(r, name):
58                return getattr(r, name)
59        raise AttributeError(name)
60
61
62class RedirectToPyDBIoMessages(object):
63
64    def __init__(self, out_ctx, wrap_stream, wrap_buffer, on_write=None):
65        '''
66        :param out_ctx:
67            1=stdout and 2=stderr
68
69        :param wrap_stream:
70            Either sys.stdout or sys.stderr.
71
72        :param bool wrap_buffer:
73            If True the buffer attribute (which wraps writing bytes) should be
74            wrapped.
75
76        :param callable(str) on_write:
77            May be a custom callable to be called when to write something.
78            If not passed the default implementation will create an io message
79            and send it through the debugger.
80        '''
81        encoding = getattr(wrap_stream, 'encoding', None)
82        if not encoding:
83            encoding = os.environ.get('PYTHONIOENCODING', 'utf-8')
84        self.encoding = encoding
85        self._out_ctx = out_ctx
86        if wrap_buffer:
87            self.buffer = RedirectToPyDBIoMessages(out_ctx, wrap_stream, wrap_buffer=False, on_write=on_write)
88        self._on_write = on_write
89
90    def get_pydb(self):
91        # Note: separate method for mocking on tests.
92        return get_global_debugger()
93
94    def flush(self):
95        pass  # no-op here
96
97    def write(self, s):
98        if self._on_write is not None:
99            self._on_write(s)
100            return
101
102        if s:
103            if IS_PY2:
104                # Need s in utf-8 bytes
105                if isinstance(s, unicode):  # noqa
106                    # Note: python 2.6 does not accept the "errors" keyword.
107                    s = s.encode('utf-8', 'replace')
108                else:
109                    s = s.decode(self.encoding, 'replace').encode('utf-8', 'replace')
110
111            else:
112                # Need s in str
113                if isinstance(s, bytes):
114                    s = s.decode(self.encoding, errors='replace')
115
116            py_db = self.get_pydb()
117            if py_db is not None:
118                # Note that the actual message contents will be a xml with utf-8, although
119                # the entry is str on py3 and bytes on py2.
120                cmd = py_db.cmd_factory.make_io_message(s, self._out_ctx)
121                if py_db.writer is not None:
122                    py_db.writer.add_command(cmd)
123
124
125class IOBuf:
126    '''This class works as a replacement for stdio and stderr.
127    It is a buffer and when its contents are requested, it will erase what
128    it has so far so that the next return will not return the same contents again.
129    '''
130
131    def __init__(self):
132        self.buflist = []
133        import os
134        self.encoding = os.environ.get('PYTHONIOENCODING', 'utf-8')
135
136    def getvalue(self):
137        b = self.buflist
138        self.buflist = []  # clear it
139        return ''.join(b)  # bytes on py2, str on py3.
140
141    def write(self, s):
142        if IS_PY2:
143            if isinstance(s, unicode):
144                # can't use 'errors' as kwargs in py 2.6
145                s = s.encode(self.encoding, 'replace')
146        else:
147            if isinstance(s, bytes):
148                s = s.decode(self.encoding, errors='replace')
149        self.buflist.append(s)
150
151    def isatty(self):
152        return False
153
154    def flush(self):
155        pass
156
157    def empty(self):
158        return len(self.buflist) == 0
159
160
161class _RedirectInfo(object):
162
163    def __init__(self, original, redirect_to):
164        self.original = original
165        self.redirect_to = redirect_to
166
167
168class _RedirectionsHolder:
169    _lock = ForkSafeLock(rlock=True)
170    _stack_stdout = []
171    _stack_stderr = []
172
173    _pydevd_stdout_redirect_ = None
174    _pydevd_stderr_redirect_ = None
175
176
177def start_redirect(keep_original_redirection=False, std='stdout', redirect_to=None):
178    '''
179    @param std: 'stdout', 'stderr', or 'both'
180    '''
181    with _RedirectionsHolder._lock:
182        if redirect_to is None:
183            redirect_to = IOBuf()
184
185        if std == 'both':
186            config_stds = ['stdout', 'stderr']
187        else:
188            config_stds = [std]
189
190        for std in config_stds:
191            original = getattr(sys, std)
192            stack = getattr(_RedirectionsHolder, '_stack_%s' % std)
193
194            if keep_original_redirection:
195                wrap_buffer = True if not IS_PY2 and hasattr(redirect_to, 'buffer') else False
196                new_std_instance = IORedirector(getattr(sys, std), redirect_to, wrap_buffer=wrap_buffer)
197                setattr(sys, std, new_std_instance)
198            else:
199                new_std_instance = redirect_to
200                setattr(sys, std, redirect_to)
201
202            stack.append(_RedirectInfo(original, new_std_instance))
203
204        return redirect_to
205
206
207def end_redirect(std='stdout'):
208    with _RedirectionsHolder._lock:
209        if std == 'both':
210            config_stds = ['stdout', 'stderr']
211        else:
212            config_stds = [std]
213        for std in config_stds:
214            stack = getattr(_RedirectionsHolder, '_stack_%s' % std)
215            redirect_info = stack.pop()
216            setattr(sys, std, redirect_info.original)
217
218
219def redirect_stream_to_pydb_io_messages(std):
220    '''
221    :param std:
222        'stdout' or 'stderr'
223    '''
224    with _RedirectionsHolder._lock:
225        redirect_to_name = '_pydevd_%s_redirect_' % (std,)
226        if getattr(_RedirectionsHolder, redirect_to_name) is None:
227            wrap_buffer = True if not IS_PY2 else False
228            original = getattr(sys, std)
229
230            redirect_to = RedirectToPyDBIoMessages(1 if std == 'stdout' else 2, original, wrap_buffer)
231            start_redirect(keep_original_redirection=True, std=std, redirect_to=redirect_to)
232
233            stack = getattr(_RedirectionsHolder, '_stack_%s' % std)
234            setattr(_RedirectionsHolder, redirect_to_name, stack[-1])
235            return True
236
237        return False
238
239
240def stop_redirect_stream_to_pydb_io_messages(std):
241    '''
242    :param std:
243        'stdout' or 'stderr'
244    '''
245    with _RedirectionsHolder._lock:
246        redirect_to_name = '_pydevd_%s_redirect_' % (std,)
247        redirect_info = getattr(_RedirectionsHolder, redirect_to_name)
248        if redirect_info is not None:  # :type redirect_info: _RedirectInfo
249            setattr(_RedirectionsHolder, redirect_to_name, None)
250
251            stack = getattr(_RedirectionsHolder, '_stack_%s' % std)
252            prev_info = stack.pop()
253
254            curr = getattr(sys, std)
255            if curr is redirect_info.redirect_to:
256                setattr(sys, std, redirect_info.original)
257
258
259@contextmanager
260def redirect_stream_to_pydb_io_messages_context():
261    with _RedirectionsHolder._lock:
262        redirecting = []
263        for std in ('stdout', 'stderr'):
264            if redirect_stream_to_pydb_io_messages(std):
265                redirecting.append(std)
266
267        try:
268            yield
269        finally:
270            for std in redirecting:
271                stop_redirect_stream_to_pydb_io_messages(std)
272
273