python fastapi 如何在中间件获取返回数据

656 阅读1分钟

相关代码如下:

from fastapi import FastAPI, Request, Response
app = FastAPI()

class aiwrap:
    def __init__(self, obj):
        self._it = iter(obj)
    def __aiter__(self):
        return self
    async def __anext__(self):
        try:
            value = next(self._it)
        except StopIteration:
            raise StopAsyncIteration
        return value
        
        
@ app.middleware("http")
async def add_process_time_header(request: Request,  call_next):
    response = await call_next(request)
    resp_body = [section async for section in response.__dict__['body_iterator']]
    response.__setattr__('body_iterator', aiwrap(resp_body))
    try:
        resp_body = json.loads(resp_body[0].decode())
    except Exception:
        resp_body = str(resp_body)
    print(resp_body)