codekingpro/portable-devtools

1"""2Application Profiler3====================4 5This module provides a middleware that profiles each request with the6:mod:`cProfile` module. This can help identify bottlenecks in your code7that may be slowing down your application.8 9.. autoclass:: ProfilerMiddleware10 11:copyright: 2007 Pallets12:license: BSD-3-Clause13"""14 15from __future__ import annotations16 17import os.path18import sys19import time20import typing as t21from pstats import Stats22 23try:24    from cProfile import Profile25except ImportError:26    from profile import Profile  # type: ignore27 28if t.TYPE_CHECKING:29    from _typeshed.wsgi import StartResponse30    from _typeshed.wsgi import WSGIApplication31    from _typeshed.wsgi import WSGIEnvironment32 33 34class ProfilerMiddleware:35    """Wrap a WSGI application and profile the execution of each36    request. Responses are buffered so that timings are more exact.37 38    If ``stream`` is given, :class:`pstats.Stats` are written to it39    after each request. If ``profile_dir`` is given, :mod:`cProfile`40    data files are saved to that directory, one file per request.41 42    The filename can be customized by passing ``filename_format``. If43    it is a string, it will be formatted using :meth:`str.format` with44    the following fields available:45 46    -   ``{method}`` - The request method; GET, POST, etc.47    -   ``{path}`` - The request path or 'root' should one not exist.48    -   ``{elapsed}`` - The elapsed time of the request in milliseconds.49    -   ``{time}`` - The time of the request.50 51    If it is a callable, it will be called with the WSGI ``environ`` and52    be expected to return a filename string. The ``environ`` dictionary53    will also have the ``"werkzeug.profiler"`` key populated with a54    dictionary containing the following fields (more may be added in the55    future):56    -   ``{elapsed}`` - The elapsed time of the request in milliseconds.57    -   ``{time}`` - The time of the request.58 59    :param app: The WSGI application to wrap.60    :param stream: Write stats to this stream. Disable with ``None``.61    :param sort_by: A tuple of columns to sort stats by. See62        :meth:`pstats.Stats.sort_stats`.63    :param restrictions: A tuple of restrictions to filter stats by. See64        :meth:`pstats.Stats.print_stats`.65    :param profile_dir: Save profile data files to this directory.66    :param filename_format: Format string for profile data file names,67        or a callable returning a name. See explanation above.68 69    .. code-block:: python70 71        from werkzeug.middleware.profiler import ProfilerMiddleware72        app = ProfilerMiddleware(app)73 74    .. versionchanged:: 3.075        Added the ``"werkzeug.profiler"`` key to the ``filename_format(environ)``76        parameter with the  ``elapsed`` and ``time`` fields.77 78    .. versionchanged:: 0.1579        Stats are written even if ``profile_dir`` is given, and can be80        disable by passing ``stream=None``.81 82    .. versionadded:: 0.1583        Added ``filename_format``.84 85    .. versionadded:: 0.986        Added ``restrictions`` and ``profile_dir``.87    """88 89    def __init__(90        self,91        app: WSGIApplication,92        stream: t.IO[str] | None = sys.stdout,93        sort_by: t.Iterable[str] = ("time", "calls"),94        restrictions: t.Iterable[str | int | float] = (),95        profile_dir: str | None = None,96        filename_format: str = "{method}.{path}.{elapsed:.0f}ms.{time:.0f}.prof",97    ) -> None:98        self._app = app99        self._stream = stream100        self._sort_by = sort_by101        self._restrictions = restrictions102        self._profile_dir = profile_dir103        self._filename_format = filename_format104 105    def __call__(106        self, environ: WSGIEnvironment, start_response: StartResponse107    ) -> t.Iterable[bytes]:108        response_body: list[bytes] = []109 110        def catching_start_response(status, headers, exc_info=None):  # type: ignore111            start_response(status, headers, exc_info)112            return response_body.append113 114        def runapp() -> None:115            app_iter = self._app(116                environ, t.cast("StartResponse", catching_start_response)117            )118            response_body.extend(app_iter)119 120            if hasattr(app_iter, "close"):121                app_iter.close()122 123        profile = Profile()124        start = time.time()125        profile.runcall(runapp)126        body = b"".join(response_body)127        elapsed = time.time() - start128 129        if self._profile_dir is not None:130            if callable(self._filename_format):131                environ["werkzeug.profiler"] = {132                    "elapsed": elapsed * 1000.0,133                    "time": time.time(),134                }135                filename = self._filename_format(environ)136            else:137                filename = self._filename_format.format(138                    method=environ["REQUEST_METHOD"],139                    path=environ["PATH_INFO"].strip("/").replace("/", ".") or "root",140                    elapsed=elapsed * 1000.0,141                    time=time.time(),142                )143            filename = os.path.join(self._profile_dir, filename)144            profile.dump_stats(filename)145 146        if self._stream is not None:147            stats = Stats(profile, stream=self._stream)148            stats.sort_stats(*self._sort_by)149            print("-" * 80, file=self._stream)150            path_info = environ.get("PATH_INFO", "")151            print(f"PATH: {path_info!r}", file=self._stream)152            stats.print_stats(*self._restrictions)153            print(f"{'-' * 80}\n", file=self._stream)154 155        return [body]156