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"""
4Middleware that profiles the request and displays profiling
5information at the bottom of each page.
6"""
7
8
9import sys
10import os
11import hotshot
12import hotshot.stats
13import threading
14import cgi
15import six
16import time
17from cStringIO import StringIO
18from paste import response
19
20__all__ = ['ProfileMiddleware', 'profile_decorator']
21
22class ProfileMiddleware(object):
23
24    """
25    Middleware that profiles all requests.
26
27    All HTML pages will have profiling information appended to them.
28    The data is isolated to that single request, and does not include
29    data from previous requests.
30
31    This uses the ``hotshot`` module, which affects performance of the
32    application.  It also runs in a single-threaded mode, so it is
33    only usable in development environments.
34    """
35
36    style = ('clear: both; background-color: #ff9; color: #000; '
37             'border: 2px solid #000; padding: 5px;')
38
39    def __init__(self, app, global_conf=None,
40                 log_filename='profile.log.tmp',
41                 limit=40):
42        self.app = app
43        self.lock = threading.Lock()
44        self.log_filename = log_filename
45        self.limit = limit
46
47    def __call__(self, environ, start_response):
48        catch_response = []
49        body = []
50        def replace_start_response(status, headers, exc_info=None):
51            catch_response.extend([status, headers])
52            start_response(status, headers, exc_info)
53            return body.append
54        def run_app():
55            app_iter = self.app(environ, replace_start_response)
56            try:
57                body.extend(app_iter)
58            finally:
59                if hasattr(app_iter, 'close'):
60                    app_iter.close()
61        self.lock.acquire()
62        try:
63            prof = hotshot.Profile(self.log_filename)
64            prof.addinfo('URL', environ.get('PATH_INFO', ''))
65            try:
66                prof.runcall(run_app)
67            finally:
68                prof.close()
69            body = ''.join(body)
70            headers = catch_response[1]
71            content_type = response.header_value(headers, 'content-type')
72            if content_type is None or not content_type.startswith('text/html'):
73                # We can't add info to non-HTML output
74                return [body]
75            stats = hotshot.stats.load(self.log_filename)
76            stats.strip_dirs()
77            stats.sort_stats('time', 'calls')
78            output = capture_output(stats.print_stats, self.limit)
79            output_callers = capture_output(
80                stats.print_callers, self.limit)
81            body += '<pre style="%s">%s\n%s</pre>' % (
82                self.style, cgi.escape(output), cgi.escape(output_callers))
83            return [body]
84        finally:
85            self.lock.release()
86
87def capture_output(func, *args, **kw):
88    # Not threadsafe! (that's okay when ProfileMiddleware uses it,
89    # though, since it synchronizes itself.)
90    out = StringIO()
91    old_stdout = sys.stdout
92    sys.stdout = out
93    try:
94        func(*args, **kw)
95    finally:
96        sys.stdout = old_stdout
97    return out.getvalue()
98
99def profile_decorator(**options):
100
101    """
102    Profile a single function call.
103
104    Used around a function, like::
105
106        @profile_decorator(options...)
107        def ...
108
109    All calls to the function will be profiled.  The options are
110    all keywords, and are:
111
112        log_file:
113            The filename to log to (or ``'stdout'`` or ``'stderr'``).
114            Default: stderr.
115        display_limit:
116            Only show the top N items, default: 20.
117        sort_stats:
118            A list of string-attributes to sort on.  Default
119            ``('time', 'calls')``.
120        strip_dirs:
121            Strip directories/module names from files?  Default True.
122        add_info:
123            If given, this info will be added to the report (for your
124            own tracking).  Default: none.
125        log_filename:
126            The temporary filename to log profiling data to.  Default;
127            ``./profile_data.log.tmp``
128        no_profile:
129            If true, then don't actually profile anything.  Useful for
130            conditional profiling.
131    """
132
133    if options.get('no_profile'):
134        def decorator(func):
135            return func
136        return decorator
137    def decorator(func):
138        def replacement(*args, **kw):
139            return DecoratedProfile(func, **options)(*args, **kw)
140        return replacement
141    return decorator
142
143class DecoratedProfile(object):
144
145    lock = threading.Lock()
146
147    def __init__(self, func, **options):
148        self.func = func
149        self.options = options
150
151    def __call__(self, *args, **kw):
152        self.lock.acquire()
153        try:
154            return self.profile(self.func, *args, **kw)
155        finally:
156            self.lock.release()
157
158    def profile(self, func, *args, **kw):
159        ops = self.options
160        prof_filename = ops.get('log_filename', 'profile_data.log.tmp')
161        prof = hotshot.Profile(prof_filename)
162        prof.addinfo('Function Call',
163                     self.format_function(func, *args, **kw))
164        if ops.get('add_info'):
165            prof.addinfo('Extra info', ops['add_info'])
166        exc_info = None
167        try:
168            start_time = time.time()
169            try:
170                result = prof.runcall(func, *args, **kw)
171            except:
172                exc_info = sys.exc_info()
173            end_time = time.time()
174        finally:
175            prof.close()
176        stats = hotshot.stats.load(prof_filename)
177        os.unlink(prof_filename)
178        if ops.get('strip_dirs', True):
179            stats.strip_dirs()
180        stats.sort_stats(*ops.get('sort_stats', ('time', 'calls')))
181        display_limit = ops.get('display_limit', 20)
182        output = capture_output(stats.print_stats, display_limit)
183        output_callers = capture_output(
184            stats.print_callers, display_limit)
185        output_file = ops.get('log_file')
186        if output_file in (None, 'stderr'):
187            f = sys.stderr
188        elif output_file in ('-', 'stdout'):
189            f = sys.stdout
190        else:
191            f = open(output_file, 'a')
192            f.write('\n%s\n' % ('-'*60))
193            f.write('Date: %s\n' % time.strftime('%c'))
194        f.write('Function call: %s\n'
195                % self.format_function(func, *args, **kw))
196        f.write('Wall time: %0.2f seconds\n'
197                % (end_time - start_time))
198        f.write(output)
199        f.write(output_callers)
200        if output_file not in (None, '-', 'stdout', 'stderr'):
201            f.close()
202        if exc_info:
203            # We captured an exception earlier, now we re-raise it
204            six.reraise(exc_info[0], exc_info[1], exc_info[2])
205        return result
206
207    def format_function(self, func, *args, **kw):
208        args = map(repr, args)
209        args.extend(
210            ['%s=%r' % (k, v) for k, v in kw.items()])
211        return '%s(%s)' % (func.__name__, ', '.join(args))
212
213
214def make_profile_middleware(
215    app, global_conf,
216    log_filename='profile.log.tmp',
217    limit=40):
218    """
219    Wrap the application in a component that will profile each
220    request.  The profiling data is then appended to the output
221    of each page.
222
223    Note that this serializes all requests (i.e., removing
224    concurrency).  Therefore never use this in production.
225    """
226    limit = int(limit)
227    return ProfileMiddleware(
228        app, log_filename=log_filename, limit=limit)
229