Skip to content

Commit

Permalink
Include websockets
Browse files Browse the repository at this point in the history
  • Loading branch information
davidbrochart committed Dec 22, 2021
1 parent 57c2db2 commit 842b7eb
Showing 1 changed file with 21 additions and 5 deletions.
26 changes: 21 additions & 5 deletions fps/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

import pluggy
from fastapi import FastAPI
from fastapi.routing import APIWebSocketRoute
from pluggy import PluginManager
from rich.console import Console
from rich.table import Table
Expand Down Expand Up @@ -155,7 +156,7 @@ def _load_configurations() -> None:
logger.info("No plugin configuration to load")


def _load_routers(app: FastAPI) -> None:
def _load_routers(app: FastAPI) -> Dict[str, APIWebSocketRoute]:

pm = _get_pluggin_manager(HookType.ROUTER)

Expand All @@ -176,6 +177,7 @@ def _load_routers(app: FastAPI) -> None:
registered_paths = {}
registered_routes = {}
registered_mounts = {}
ws_routes = {}

for p, routers in grouped_routers.items():
p_name = Config.plugin_name(p)
Expand Down Expand Up @@ -259,6 +261,13 @@ def _load_routers(app: FastAPI) -> None:
tags=tags,
)
routes_count += len(routes)
ws_routes.update(
{
tags[0]: route
for route in routes
if isinstance(route, APIWebSocketRoute)
}
)

if router_mounts:
for m in router_mounts.values():
Expand All @@ -272,20 +281,27 @@ def _load_routers(app: FastAPI) -> None:
else:
logger.info("No plugin API router to load")

return ws_routes


def show_endpoints(app: FastAPI):
def show_endpoints(app: FastAPI, ws_routes: Dict[str, APIWebSocketRoute]):
table = Table(title="API Summary")
table.add_column("Path", justify="left", style="cyan", no_wrap=True)
table.add_column("Methods", justify="right", style="green")
table.add_column("Plugin", style="magenta")

# HTTP endpoints
openapi = app.openapi()
for k, v in openapi["paths"].items():
path = k
methods = ", ".join([method.upper() for method in v.keys()])
plugin = ", ".join({tag for i in v.values() for tag in i["tags"]})
plugin = ", ".join({i["tags"][0] for i in v.values()})
table.add_row(path, methods, plugin)

# websockets endpoints
for plugin, ws_route in ws_routes.items():
table.add_row(f"[cyan on red]{ws_route.path}[/]", "WEBSOCKET", plugin)

console = Console()
with console.capture() as capture:
console.print()
Expand All @@ -305,11 +321,11 @@ def create_app():
fps_config = Config(FPSConfig)
app = FastAPI(**fps_config.__dict__)

_load_routers(app)
ws_routes = _load_routers(app)
_load_exceptions_handlers(app)

Config.check_not_used_sections()

show_endpoints(app)
show_endpoints(app, ws_routes)

return app

0 comments on commit 842b7eb

Please sign in to comment.