pycaptcha/build_app.py

import os

from fastapi import FastAPI, Request
from fastapi.openapi.docs import (
    get_redoc_html,
    get_swagger_ui_html,
    get_swagger_ui_oauth2_redirect_html,
)
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles

from exception import ApiException
from services.limiter import register_limiter
from services.routers import service_router


def register_swagger_ui_html(app):
    app.mount("/static", StaticFiles(directory="static"), name="static")

    @app.get("/docs", include_in_schema=False)
    async def custom_swagger_ui_html():
        return get_swagger_ui_html(
            openapi_url=app.openapi_url,
            title=app.title + " - Swagger UI",
            oauth2_redirect_url=app.swagger_ui_oauth2_redirect_url,
            swagger_js_url="/static/swagger-ui-bundle.js",
            swagger_css_url="/static/swagger-ui.css",
        )

    @app.get(app.swagger_ui_oauth2_redirect_url, include_in_schema=False)
    async def swagger_ui_redirect():
        return get_swagger_ui_oauth2_redirect_html()

    @app.get("/redoc", include_in_schema=False)
    async def redoc_html():
        return get_redoc_html(
            openapi_url=app.openapi_url,
            title=app.title + " - ReDoc",
            redoc_js_url="/static/redoc.standalone.js",
        )


def init_app():
    env = os.getenv("FASTAPI_ENV")
    if env != "prod":
        app = FastAPI(title="Captcha", docs_url=None, redoc_url=None)
        register_swagger_ui_html(app)
    else:
        # 生产环境关闭文档说明
        app = FastAPI(title="Captcha", docs_url=None, redoc_url=None)

    app.include_router(service_router, prefix="/v1")

    @app.exception_handler(ApiException)
    async def api_exception(request: Request, exc: ApiException):
        return JSONResponse(status_code=200, content=dict(code=exc.code, errMsg=exc.errMsg, r=exc.r))

    @app.exception_handler(Exception)
    async def unknown_exception(request: Request, exc: Exception):
        return JSONResponse(status_code=500, content=dict(code=1, errMsg="服务器内部错误,暂无法提供服务", r={}))

    return app


def create_app():
    app = init_app()
    register_limiter(app)
    return app