How to close a socket inside an fastapi StreamingResponse generator?-CodePudding

The code below doesn't reach out of the with. Is there any way I can close the socket when the client disconnects?

import fastapi
import uvicorn
import socket

def gen_video(port):
    with socket.socket() as s:
        s.connect(("127.0.0.1", port))
        while True:
            yield s.recv(1024)


app = fastapi.FastAPI()

@app.get("/stream")
async def stream():
    return fastapi.responses.StreamingResponse(gen_video(1234), media_type="video/ogg")

if __name__ == "__main__":
    uvicorn.run(app, port=8000)

CodePudding user response：

An approach using fastapi BackgroundTasks to close the socket

import fastapi
import uvicorn
import socket

def read_socket(s):
    while True:
        yield s.recv(1024)

app = fastapi.FastAPI()

@app.get("/stream")
async def stream(background_tasks: fastapi.background.BackgroundTasks):
    s = socket.socket()
    s.connect(("127.0.0.1", 1234))
    background_tasks.add_task(lambda f: f.close(), s)
    return fastapi.responses.StreamingResponse(read_socket(s), media_type="video/ogg")

if __name__ == "__main__":
    uvicorn.run(app, port=8000)