Merge pull request #11 from rgeoghegan/fix-html-mode

Fix html mode; move PyInstrumentProfilerMiddleware to top-level module

example/fastapi_example.py

@@ -4,7 +4,7 @@
 from fastapi import FastAPI
 from fastapi.responses import JSONResponse
 
-from fastapi_profiler.profiler_middleware import PyInstrumentProfilerMiddleware
+from fastapi_profiler import PyInstrumentProfilerMiddleware
 
 
 app = FastAPI()

example/fastapi_to_html_example.py

@@ -0,0 +1,36 @@
+"""
+This example shows how to output the profile
+to an html file.
+"""
+import os
+import uvicorn
+
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse
+
+from fastapi_profiler import PyInstrumentProfilerMiddleware
+
+
+app = FastAPI()
+app.add_middleware(
+    PyInstrumentProfilerMiddleware,
+    server_app=app,  # Required to output the profile on server shutdown
+    profiler_output_type="html",
+    is_print_each_request=False,  # Set to True to show request profile on
+                                  # stdout on each request
+    open_in_browser=False,  # Set to true to open your web-browser automatically
+                            # when the server shuts down
+    html_file_name="example_profile.html"  # Filename for output
+)
+
+
+@app.get("/test")
+async def normal_request():
+    return JSONResponse({"retMsg": "Hello World!"})
+
+
+# Or you can use the console with command "uvicorn" to run this example.
+# Command: uvicorn fastapi_example:app --host="0.0.0.0" --port=8080
+if __name__ == '__main__':
+    app_name = os.path.basename(__file__).replace(".py", "")
+    uvicorn.run(app=f"{app_name}:app", host="0.0.0.0", port=8080, workers=1)

fastapi_profiler/__init__.py

@@ -1,2 +1,2 @@
-from fastapi_profiler import profiler_middleware
-from fastapi_profiler._version import __author__, __version__
+from ._version import __version__, __author__
+from .profiler import PyInstrumentProfilerMiddleware

fastapi_profiler/_version.py

@@ -1,2 +1,2 @@
-__version__ = "1.0.0"
+__version__ = "2.0.0"
 __author__ = "sunhailin-Leo"

fastapi_profiler/profiler_middleware/profiler.py → fastapi_profiler/profiler.py

@@ -1,3 +1,7 @@
+__version__ = "1.0.0"
+__author__ = "sunhailin-Leo"
+
+import os
 import time
 import codecs
 from typing import Optional
@@ -15,28 +19,38 @@
 
 
 class PyInstrumentProfilerMiddleware:
+    DEFAULT_HTML_FILENAME = "./fastapi-profiler.html"
+
     def __init__(
         self, app: ASGIApp,
         *,
         server_app: Optional[Router] = None,
         profiler_interval: float = 0.0001,
         profiler_output_type: str = "text",
         is_print_each_request: bool = True,
+        html_file_name: Optional[str] = None,
+        open_in_browser: bool = False,
         **profiler_kwargs
     ):
         self.app = app
         self._profiler = Profiler(interval=profiler_interval)
-
-        self._server_app = server_app
         self._output_type = profiler_output_type
         self._print_each_request = is_print_each_request
+        self._html_file_name: Optional[str] = html_file_name
+        self._open_in_browser: bool = open_in_browser
         self._profiler_kwargs: dict = profiler_kwargs
 
-    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if profiler_output_type == "html" and server_app is None:
+            raise RuntimeError(
+                "If profiler_output_type=html, must provide server_app argument "
+                "to set shutdown event handler to output profile."
+            )
+
         # register an event handler for profiler stop
-        if self._server_app is not None:
-            self._server_app.add_event_handler("shutdown", self.get_profiler_result)
+        if server_app is not None:
+            server_app.add_event_handler("shutdown", self.get_profiler_result)
 
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if scope["type"] != "http":
             await self.app(scope, receive, send)
             return
@@ -73,33 +87,27 @@ async def wrapped_send(message: Message) -> None:
 
     async def get_profiler_result(self):
         if self._output_type == "text":
+            logger.info("Compiling and printing final profile")
             print(self._profiler.output_text(**self._profiler_kwargs))
         elif self._output_type == "html":
-            html_name = self._profiler_kwargs.get("html_file_name")
-            if html_name is None:
-                html_name = "fastapi-profiler.html"
-
-            """
-             There are some problems with the args -- output_filename.
-             You can check the
-                class
-                    'from pyinstrument.renderers import HTMLRenderer'
-                method
-                    'open_in_browser'
-             the argument 'output_filename' will become the URL like 'file://xxxx',
-             but that code have some bugs on it.
-
-             So on my middleware, the args 'html_file_name'
-             I suggest use None to instead, or you can use the absolute path.
-
-             HTMLRenderer().open_in_browser(
-                session=self._profiler.last_session,
-                output_filename=html_name,
-             )
-
-             At last, I rewrite the function to avoid the problem!
-             By the way, the html file default save at the root path of your project.
-            """
-            html_code = HTMLRenderer().render(session=self._profiler.last_session)
-            with codecs.open(html_name, "w", "utf-8") as f:
-                f.write(html_code)
+            html_file_name = self.DEFAULT_HTML_FILENAME
+            if self._html_file_name is not None:
+                html_file_name = self._html_file_name
+
+            logger.info(
+                "Compiling and dumping final profile to %r - this may take some time",
+                html_file_name,
+            )
+
+            renderer = HTMLRenderer()
+            if self._open_in_browser:
+                renderer.open_in_browser(
+                    session=self._profiler.last_session,
+                    output_filename=os.path.abspath(html_file_name),
+                )
+            else:
+                html_code = renderer.render(session=self._profiler.last_session)
+                with codecs.open(html_file_name, "w", "utf-8") as f:
+                    f.write(html_code)
+
+            logger.info("Done writing profile to %r", html_file_name)

test/test_middleware.py

@@ -8,7 +8,7 @@
 from fastapi.responses import JSONResponse
 
 from test import stdout_redirect
-from fastapi_profiler.profiler_middleware import PyInstrumentProfilerMiddleware
+from fastapi_profiler import PyInstrumentProfilerMiddleware
 
 
 @pytest.fixture(name="test_middleware")
@@ -45,17 +45,17 @@ def test_profiler_print_at_console(self, client):
         sys.stdout = temp_stdout
         assert (f"Path: {request_path}" in stdout_redirect.fp.getvalue())
 
-    def test_profiler_export_to_html(self, test_middleware):
-        full_path = f"{os.getcwd()}/test.html"
+    def test_profiler_export_to_html(self, test_middleware, tmpdir):
+        full_path = tmpdir / "test.html"
 
         with TestClient(test_middleware(
                 profiler_output_type="html",
                 is_print_each_request=False,
-                html_file_name=full_path)) as client:
+                profiler_interval=0.0000001,
+                html_file_name=str(full_path))) as client:
             # request
             request_path = "/test"
             client.get(request_path)
 
         # HTML will record the py file name.
-        with open(full_path, "r") as f:
-            assert ("profiler.py" in f.read())
+        assert "profiler.py" in full_path.read_text("utf-8")