1# (c) 2005 Ian Bicking and contributors; written for Paste (http://pythonpaste.org)
2# Licensed under the MIT license: http://www.opensource.org/licenses/mit-license.php
3
4"""
5threadedprint.py
6================
7
8:author: Ian Bicking
9:date: 12 Jul 2004
10
11Multi-threaded printing; allows the output produced via print to be
12separated according to the thread.
13
14To use this, you must install the catcher, like::
15
16    threadedprint.install()
17
18The installation optionally takes one of three parameters:
19
20default
21    The default destination for print statements (e.g., ``sys.stdout``).
22factory
23    A function that will produce the stream for a thread, given the
24    thread's name.
25paramwriter
26    Instead of writing to a file-like stream, this function will be
27    called like ``paramwriter(thread_name, text)`` for every write.
28
29The thread name is the value returned by
30``threading.currentThread().getName()``, a string (typically something
31like Thread-N).
32
33You can also submit file-like objects for specific threads, which will
34override any of these parameters.  To do this, call ``register(stream,
35[threadName])``.  ``threadName`` is optional, and if not provided the
36stream will be registered for the current thread.
37
38If no specific stream is registered for a thread, and no default has
39been provided, then an error will occur when anything is written to
40``sys.stdout`` (or printed).
41
42Note: the stream's ``write`` method will be called in the thread the
43text came from, so you should consider thread safety, especially if
44multiple threads share the same writer.
45
46Note: if you want access to the original standard out, use
47``sys.__stdout__``.
48
49You may also uninstall this, via::
50
51    threadedprint.uninstall()
52
53TODO
54----
55
56* Something with ``sys.stderr``.
57* Some default handlers.  Maybe something that hooks into `logging`.
58* Possibly cache the results of ``factory`` calls.  This would be a
59  semantic change.
60
61"""
62
63import threading
64import sys
65from paste.util import filemixin
66
67class PrintCatcher(filemixin.FileMixin):
68
69    def __init__(self, default=None, factory=None, paramwriter=None,
70                 leave_stdout=False):
71        assert len(filter(lambda x: x is not None,
72                          [default, factory, paramwriter])) <= 1, (
73            "You can only provide one of default, factory, or paramwriter")
74        if leave_stdout:
75            assert not default, (
76                "You cannot pass in both default (%r) and "
77                "leave_stdout=True" % default)
78            default = sys.stdout
79        if default:
80            self._defaultfunc = self._writedefault
81        elif factory:
82            self._defaultfunc = self._writefactory
83        elif paramwriter:
84            self._defaultfunc = self._writeparam
85        else:
86            self._defaultfunc = self._writeerror
87        self._default = default
88        self._factory = factory
89        self._paramwriter = paramwriter
90        self._catchers = {}
91
92    def write(self, v, currentThread=threading.currentThread):
93        name = currentThread().getName()
94        catchers = self._catchers
95        if not catchers.has_key(name):
96            self._defaultfunc(name, v)
97        else:
98            catcher = catchers[name]
99            catcher.write(v)
100
101    def seek(self, *args):
102        # Weird, but Google App Engine is seeking on stdout
103        name = threading.currentThread().getName()
104        catchers = self._catchers
105        if not name in catchers:
106            self._default.seek(*args)
107        else:
108            catchers[name].seek(*args)
109
110    def read(self, *args):
111        name = threading.currentThread().getName()
112        catchers = self._catchers
113        if not name in catchers:
114            self._default.read(*args)
115        else:
116            catchers[name].read(*args)
117
118
119    def _writedefault(self, name, v):
120        self._default.write(v)
121
122    def _writefactory(self, name, v):
123        self._factory(name).write(v)
124
125    def _writeparam(self, name, v):
126        self._paramwriter(name, v)
127
128    def _writeerror(self, name, v):
129        assert False, (
130            "There is no PrintCatcher output stream for the thread %r"
131            % name)
132
133    def register(self, catcher, name=None,
134                 currentThread=threading.currentThread):
135        if name is None:
136            name = currentThread().getName()
137        self._catchers[name] = catcher
138
139    def deregister(self, name=None,
140                   currentThread=threading.currentThread):
141        if name is None:
142            name = currentThread().getName()
143        assert self._catchers.has_key(name), (
144            "There is no PrintCatcher catcher for the thread %r" % name)
145        del self._catchers[name]
146
147_printcatcher = None
148_oldstdout = None
149
150def install(**kw):
151    global _printcatcher, _oldstdout, register, deregister
152    if (not _printcatcher or sys.stdout is not _printcatcher):
153        _oldstdout = sys.stdout
154        _printcatcher = sys.stdout = PrintCatcher(**kw)
155        register = _printcatcher.register
156        deregister = _printcatcher.deregister
157
158def uninstall():
159    global _printcatcher, _oldstdout, register, deregister
160    if _printcatcher:
161        sys.stdout = _oldstdout
162        _printcatcher = _oldstdout = None
163        register = not_installed_error
164        deregister = not_installed_error
165
166def not_installed_error(*args, **kw):
167    assert False, (
168        "threadedprint has not yet been installed (call "
169        "threadedprint.install())")
170
171register = deregister = not_installed_error
172
173class StdinCatcher(filemixin.FileMixin):
174
175    def __init__(self, default=None, factory=None, paramwriter=None):
176        assert len(filter(lambda x: x is not None,
177                          [default, factory, paramwriter])) <= 1, (
178            "You can only provide one of default, factory, or paramwriter")
179        if default:
180            self._defaultfunc = self._readdefault
181        elif factory:
182            self._defaultfunc = self._readfactory
183        elif paramwriter:
184            self._defaultfunc = self._readparam
185        else:
186            self._defaultfunc = self._readerror
187        self._default = default
188        self._factory = factory
189        self._paramwriter = paramwriter
190        self._catchers = {}
191
192    def read(self, size=None, currentThread=threading.currentThread):
193        name = currentThread().getName()
194        catchers = self._catchers
195        if not catchers.has_key(name):
196            return self._defaultfunc(name, size)
197        else:
198            catcher = catchers[name]
199            return catcher.read(size)
200
201    def _readdefault(self, name, size):
202        self._default.read(size)
203
204    def _readfactory(self, name, size):
205        self._factory(name).read(size)
206
207    def _readparam(self, name, size):
208        self._paramreader(name, size)
209
210    def _readerror(self, name, size):
211        assert False, (
212            "There is no StdinCatcher output stream for the thread %r"
213            % name)
214
215    def register(self, catcher, name=None,
216                 currentThread=threading.currentThread):
217        if name is None:
218            name = currentThread().getName()
219        self._catchers[name] = catcher
220
221    def deregister(self, catcher, name=None,
222                   currentThread=threading.currentThread):
223        if name is None:
224            name = currentThread().getName()
225        assert self._catchers.has_key(name), (
226            "There is no StdinCatcher catcher for the thread %r" % name)
227        del self._catchers[name]
228
229_stdincatcher = None
230_oldstdin = None
231
232def install_stdin(**kw):
233    global _stdincatcher, _oldstdin, register_stdin, deregister_stdin
234    if not _stdincatcher:
235        _oldstdin = sys.stdin
236        _stdincatcher = sys.stdin = StdinCatcher(**kw)
237        register_stdin = _stdincatcher.register
238        deregister_stdin = _stdincatcher.deregister
239
240def uninstall_stdin():
241    global _stdincatcher, _oldstdin, register_stdin, deregister_stdin
242    if _stdincatcher:
243        sys.stdin = _oldstdin
244        _stdincatcher = _oldstdin = None
245        register_stdin = deregister_stdin = not_installed_error_stdin
246
247def not_installed_error_stdin(*args, **kw):
248    assert False, (
249        "threadedprint has not yet been installed for stdin (call "
250        "threadedprint.install_stdin())")
251