profiler.py 4.8 KB

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