venv/lib/python3.11/site-packages/werkzeug/middleware/profiler.py

   1 """
   2 Application Profiler
   3 ====================
   4
   5 This module provides a middleware that profiles each request with the
   6 :mod:`cProfile` module. This can help identify bottlenecks in your code
   7 that may be slowing down your application.
   8
   9 .. autoclass:: ProfilerMiddleware
  10
  11 :copyright: 2007 Pallets
  12 :license: BSD-3-Clause
  13 """
  14 from __future__ import annotations
  15
  16 import os.path
  17 import sys
  18 import time
  19 import typing as t
  20 from pstats import Stats
  21
  22 try:
  23     from cProfile import Profile
  24 except ImportError:
  25     from profile import Profile  # type: ignore
  26
  27 if t.TYPE_CHECKING:
  28     from _typeshed.wsgi import StartResponse
  29     from _typeshed.wsgi import WSGIApplication
  30     from _typeshed.wsgi import WSGIEnvironment
  31
  32
  33 class ProfilerMiddleware:
  34     """Wrap a WSGI application and profile the execution of each
  35     request. Responses are buffered so that timings are more exact.
  36
  37     If ``stream`` is given, :class:`pstats.Stats` are written to it
  38     after each request. If ``profile_dir`` is given, :mod:`cProfile`
  39     data files are saved to that directory, one file per request.
  40
  41     The filename can be customized by passing ``filename_format``. If
  42     it is a string, it will be formatted using :meth:`str.format` with
  43     the following fields available:
  44
  45     -   ``{method}`` - The request method; GET, POST, etc.
  46     -   ``{path}`` - The request path or 'root' should one not exist.
  47     -   ``{elapsed}`` - The elapsed time of the request in milliseconds.
  48     -   ``{time}`` - The time of the request.
  49
  50     If it is a callable, it will be called with the WSGI ``environ`` and
  51     be expected to return a filename string. The ``environ`` dictionary
  52     will also have the ``"werkzeug.profiler"`` key populated with a
  53     dictionary containing the following fields (more may be added in the
  54     future):
  55     -   ``{elapsed}`` - The elapsed time of the request in milliseconds.
  56     -   ``{time}`` - The time of the request.
  57
  58     :param app: The WSGI application to wrap.
  59     :param stream: Write stats to this stream. Disable with ``None``.
  60     :param sort_by: A tuple of columns to sort stats by. See
  61         :meth:`pstats.Stats.sort_stats`.
  62     :param restrictions: A tuple of restrictions to filter stats by. See
  63         :meth:`pstats.Stats.print_stats`.
  64     :param profile_dir: Save profile data files to this directory.
  65     :param filename_format: Format string for profile data file names,
  66         or a callable returning a name. See explanation above.
  67
  68     .. code-block:: python
  69
  70         from werkzeug.middleware.profiler import ProfilerMiddleware
  71         app = ProfilerMiddleware(app)
  72
  73     .. versionchanged:: 3.0
  74         Added the ``"werkzeug.profiler"`` key to the ``filename_format(environ)``
  75         parameter with the  ``elapsed`` and ``time`` fields.
  76
  77     .. versionchanged:: 0.15
  78         Stats are written even if ``profile_dir`` is given, and can be
  79         disable by passing ``stream=None``.
  80
  81     .. versionadded:: 0.15
  82         Added ``filename_format``.
  83
  84     .. versionadded:: 0.9
  85         Added ``restrictions`` and ``profile_dir``.
  86     """
  87
  88     def __init__(
  89         self,
  90         app: WSGIApplication,
  91         stream: t.IO[str] | None = sys.stdout,
  92         sort_by: t.Iterable[str] = ("time", "calls"),
  93         restrictions: t.Iterable[str | int | float] = (),
  94         profile_dir: str | None = None,
  95         filename_format: str = "{method}.{path}.{elapsed:.0f}ms.{time:.0f}.prof",
  96     ) -> None:
  97         self._app = app
  98         self._stream = stream
  99         self._sort_by = sort_by
 100         self._restrictions = restrictions
 101         self._profile_dir = profile_dir
 102         self._filename_format = filename_format
 103
 104     def __call__(
 105         self, environ: WSGIEnvironment, start_response: StartResponse
 106     ) -> t.Iterable[bytes]:
 107         response_body: list[bytes] = []
 108
 109         def catching_start_response(status, headers, exc_info=None):  # type: ignore
 110             start_response(status, headers, exc_info)
 111             return response_body.append
 112
 113         def runapp() -> None:
 114             app_iter = self._app(
 115                 environ, t.cast("StartResponse", catching_start_response)
 116             )
 117             response_body.extend(app_iter)
 118
 119             if hasattr(app_iter, "close"):
 120                 app_iter.close()
 121
 122         profile = Profile()
 123         start = time.time()
 124         profile.runcall(runapp)
 125         body = b"".join(response_body)
 126         elapsed = time.time() - start
 127
 128         if self._profile_dir is not None:
 129             if callable(self._filename_format):
 130                 environ["werkzeug.profiler"] = {
 131                     "elapsed": elapsed * 1000.0,
 132                     "time": time.time(),
 133                 }
 134                 filename = self._filename_format(environ)
 135             else:
 136                 filename = self._filename_format.format(
 137                     method=environ["REQUEST_METHOD"],
 138                     path=environ["PATH_INFO"].strip("/").replace("/", ".") or "root",
 139                     elapsed=elapsed * 1000.0,
 140                     time=time.time(),
 141                 )
 142             filename = os.path.join(self._profile_dir, filename)
 143             profile.dump_stats(filename)
 144
 145         if self._stream is not None:
 146             stats = Stats(profile, stream=self._stream)
 147             stats.sort_stats(*self._sort_by)
 148             print("-" * 80, file=self._stream)
 149             path_info = environ.get("PATH_INFO", "")
 150             print(f"PATH: {path_info!r}", file=self._stream)
 151             stats.print_stats(*self._restrictions)
 152             print(f"{'-' * 80}\n", file=self._stream)
 153
 154         return [body]