Flask-SocketIO Automated Documentation and Validation

What is Flask-SocketIO, SocketIO?

If you are reading this blog, you are probably already familiar with SocketIO python-socketio and Flask-SocketIO. If you want to learn more about them, please check the following resources:

The problem

Imagine that you are working on a large project that uses a Flask-SocketIO server to handle real-time communication between the client and the server. The server was originally well-documented, but over time the documentation has become out of date as the server has evolved and new features have been added.

I found myself in the same situation. I needed to maintain and constantly add more documentation to a Flask-SocketIO server. To make this process more efficient, I sought a solution to automate documentation generation from the existing codebase. This would eliminate the need for team members to manually edit the AsyncAPI specification file every time there was a change, a challenging task for those unfamiliar with AsyncAPI. By automating this process, we could save time and reduce the workload for the team.

To address this issue, I decided to implement SIO-AsyncAPI. This tool allows you to generate an AsyncAPI specification from your SocketIO server and validate incoming and outgoing messages against it. This functionality is similar to how FastAPI, Flask-RESTful, and other frameworks have long provided it for RESTful servers. Now, with SIO-AsyncAPI, it is possible to apply this approach to SocketIO servers as well.

How to use SIO-AsyncAPI

Instead of giving you a detailed, step-by-step guide, we'll use a more exciting approach. We'll take the existing Flask-SocketIO server and add SIO-AsyncAPI. To make things even more interesting, we'll ask ChatGPT to generate a server for us and use it as a real-world Flask-SocketIO server.

Generate a Flask-SocketIO server

I had to ask ChatGPT for multiple alterations to get the desired result. Here are my queries to ChatGPT:

create python flask-socket server that provides tic tac toe game api

change previous code and add type hints

for previous example use pydantic models instead of dictionaries

And here is the final code that ChatGPT generated for us:

1from typing import List, Union
2
3from flask import Flask
4from flask_socketio import SocketIO, emit
5from pydantic import BaseModel
6
7app = Flask(__name__)
8socketio = SocketIO(app)
9
10games: dict[int, 'Game'] = {}
11
12class Game(BaseModel):
13    board: List[List[str]]
14    turn: str
15
16class MakeMoveData(BaseModel):
17    game_id: int
18    x: int
19    y: int
20
21@socketio.on('create_game')
22def create_game() -> None:
23    # Create a new game and add it to the list of games
24    game_id = len(games) + 1
25    games[game_id] = Game(board=[['' for _ in range(3)] for _ in range(3)], turn='X')
26
27    # Send the game id to the client
28    emit('game_created', {'game_id': game_id})
29
30@socketio.on('make_move')
31def make_move(data: MakeMoveData) -> None:
32    # Get the game and make the move
33    game = games[data.game_id]
34    board = game.board
35    turn = game.turn
36    board[data.x][data.y] = turn
37
38    # Check for a win or draw
39    result = check_game_status(board)
40
41    # Update the game state and send it to the client
42    if result == 'X':
43        emit('game_won', {'winner': 'X'})
44    elif result == 'O':
45        emit('game_won', {'winner': 'O'})
46    elif result == 'draw':
47        emit('game_drawn', {})
48    else:
49        games[data.game_id] = Game(board=board, turn='O' if turn == 'X' else 'X')
50        emit('move_made', {'board': board, 'turn': game.turn})
51
52def check_game_status(board: List[List[str]]) -> Union[str, None]:
53    # Check for a win or draw
54    for i in range(3):
55        if board[i] == ['X', 'X', 'X']:
56            return 'X'
57        if board[i] == ['O', 'O', 'O']:
58            return 'O'
59    if all(board[i][i] == 'X' for i in range(3)):
60        return 'X'
61    if all(board[i][i] == 'O' for i in range(3)):
62        return 'O'
63    if all(board[i][j] != '' for i in range(3) for j in range(3)):
64        return 'draw'
65    return None
66
67
68if __name__ == '__main__':
69    socketio.run(app, debug=True)
70

It's not quite correct (e.g. data in make_move will be passed as a dictionary, not a Pydantic model), but it's good enough for our purposes.

Generate Mermaid diagram

Now let's ask ChatGPT to generate a Mermaid diagram for us as well, so we can get a better illustration of our server:

create mermaid diagram for previous example

I had to alter the given diagram a bit, but here is the final result:

Add SIO-AsyncAPI to the server

Now let's imagine that this is our old server written in Flask-SocketIO and we want to add SIO-AsyncAPI. Here is how we would do it:

Install SIO-AsyncAPI via pip install sio-asyncapi
Change import statements

1# instead of `from flask_socketio import SocketIO`
2from sio_asyncapi import AsyncAPISocketIO as SocketIO

Add an additional argument to the AsyncAPISocketIO constructor

1socketio = SocketIO(
2    app,
3    validate=True,
4    generate_docs=True,
5    version="1.0.0",
6    title="Tic-Tac-Toe API",
7    description="Tic-Tac-Toe Game API",
8    server_url="http://localhost:5000",
9    server_name="TIC_TAC_TOE_BACKEND",
10)

Tell the @socketio.on decorator to get models from the type hint.

Note: you can also pass request_model and response_model arguments to the @socketio.on decorator instead of using type hints.

@socketio.on('make_move', get_from_typehint=True)

Now type annotations will be used to generate the AsyncAPI specification and validate incoming/outgoing messages. Note that the return value from a function is not data sent by the emit function but rather the acknowledge value that the client receives.

Add an on_emit decorator to register/document a SocketIO emit event. Since we are not defining emit function ourselves but only calling it, we need to tell SIO-AsyncAPI what to expect when emit is called. E.g

1class GameCreatedData(BaseModel):
2    game_id: int
3
4@socketio.doc_emit('game_created', GameCreatedData)
5@socketio.on('create_game')
6def create_game():
7  ...
8    emit('game_created', {'game_id': game_id})

Get AsyncAPI specification

Now we can get the AsyncAPI specification by calling the socketio.asyncapi_doc.get_yaml() function. Here is what the rendered specification looks like:

Figure 1:

Validation and Error handling

SIO-AsyncAPI will automatically validate incoming and outgoing messages. If a message is invalid, it will raise one of these 3 exceptions: EmitValidationError, RequestValidationError, or ResponseValidationError.

Flask-SocketIO has the @socketio.on_error_default decorator for default error handling that we can use. E.g.:

1@socketio.on_error_default
2def default_error_handler(e: Exception):
3    """
4    Default error handler. It is called if no other error handler is defined.
5    Handles RequestValidationError, EmitValidationError, and ResponseValidationError errors.
6    """
7    if isinstance(e, RequestValidationError):
8        logger.error(f"Request validation error: {e}")
9        return {'error': str(e)}
10    elif isinstance(e, ResponseValidationError):
11        logger.critical(f"Response validation error: {e}")
12        raise e
13    if isinstance(e, EmitValidationError):
14        logger.critical(f"Emit validation error: {e}")
15        raise e
16    else:
17        logger.critical(f"Unknown error: {e}")
18        raise e

Instead of re-raising exceptions, we can return some error interpreted as an acknowledge value sent to the client. That's what we do in the example above when there is a RequestValidationError.

This is how it looks like in FireCamp if we do not provide game_id in the make_move request:

Figure 2:

Because the make_move request may return an error in the acknowledged value now, we should add a new MakeMoveAckData model and annotate the make_move function accordingly. This will automatically update the documentation in our AsyncAPI specification.

1class MakeMoveAckData(BaseModel):
2    error: Optional[str] = Field(None, description='The error message', example='Invalid move')
3
4...
5@socketio.on('make_move', get_from_typehint=True)
6def make_move(data: MakeMoveData) -> MakeMoveAckData:
7  ...

Final Code and Specification

Here is the final code of the server. I also added examples to Pydantic models to make the specification more readable.

1import pathlib
2from typing import List, Union, Optional
3
4from flask import Flask
5from flask_socketio import emit
6from pydantic import BaseModel, Field
7from loguru import logger
8
9from sio_asyncapi import AsyncAPISocketIO as SocketIO
10from sio_asyncapi import (EmitValidationError, RequestValidationError,
11                          ResponseValidationError)
12
13app = Flask(__name__)
14
15socketio = SocketIO(
16    app,
17    validate=True,
18    generate_docs=True,
19    version="1.0.0",
20    title="Tic-Tac-Toe API",
21    description="Tic-Tac-Toe Game API",
22    server_url="http://localhost:5000",
23    server_name="TIC_TAC_TOE_BACKEND",
24)
25
26games: dict[int, 'Game'] = {}
27
28
29class Game(BaseModel):
30    board: List[List[str]] = Field(..., description='The game board', example=[
31                                   ['X', 'O', ''], ['', 'X', ''], ['', '', 'O']])
32    turn: str = Field(..., description='The current turn', example='X')
33
34
35class MakeMoveData(BaseModel):
36    game_id: int = Field(..., description='The game id', example=1)
37    x: int = Field(..., description='The x coordinate', example=0)
38    y: int = Field(..., description='The y coordinate', example=0)
39
40
41class GameCreatedData(BaseModel):
42    game_id: int = Field(..., description='The game id', example=1)
43
44
45class GameWonData(BaseModel):
46    winner: str = Field(..., description='The winner', example='X')
47
48
49class GameDrawnData(BaseModel):
50    pass
51
52
53class MoveMadeData(BaseModel):
54    board: List[List[str]] = Field(..., description='The game board', example=[
55                                   ['X', 'O', ''], ['', 'X', ''], ['', '', 'O']])
56    turn: str = Field(..., description='The current turn', example='X')
57
58class MakeMoveAckData(BaseModel):
59    error: Optional[str] = Field(None, description='The error message', example='Invalid move')
60
61@socketio.doc_emit('game_created', GameCreatedData)
62@socketio.on('create_game')
63def create_game():
64    # Create a new game and add it to the list of games
65    game_id = len(games) + 1
66    games[game_id] = Game(board=[['' for _ in range(3)] for _ in range(3)], turn='X')
67
68    # Send the game id to the client
69    emit('game_created', {'game_id': game_id})
70
71@socketio.doc_emit('game_won', GameWonData)
72@socketio.doc_emit('game_drawn', GameDrawnData)
73@socketio.doc_emit('move_made', MoveMadeData)
74@socketio.on('make_move', get_from_typehint=True)
75def make_move(data: MakeMoveData) -> MakeMoveAckData:
76    # Get the game and make the move
77    logger.info(f'Making move {data}')
78    game = games[data.game_id]
79    board = game.board
80    turn = game.turn
81    board[data.x][data.y] = turn
82
83    # Check for a win or draw
84    result = check_game_status(board)
85
86    logger.info(f'Game result: {result}')
87    # Update the game state and send it to the client
88    if result == 'X':
89        emit('game_won', {'winner': 'X'})
90    elif result == 'O':
91        emit('game_won', {'winner': 'O'})
92    elif result == 'draw':
93        emit('game_drawn', {})
94    else:
95        games[data.game_id] = Game(board=board, turn='O' if turn == 'X' else 'X')
96        emit('move_made', {'board': board, 'turn': game.turn})
97
98def check_game_status(board: List[List[str]]) -> Union[str, None]:
99    # Check for a win or draw
100    for i in range(3):
101        if board[i] == ['X', 'X', 'X']:
102            return 'X'
103        if board[i] == ['O', 'O', 'O']:
104            return 'O'
105    if all(board[i][i] == 'X' for i in range(3)):
106        return 'X'
107    if all(board[i][i] == 'O' for i in range(3)):
108        return 'O'
109    if all(board[i][j] != '' for i in range(3) for j in range(3)):
110        return 'draw'
111    return None
112
113@socketio.on_error_default
114def default_error_handler(e: Exception):
115    """
116    default error handler. it called if no other error handler defined.
117    handles requestvalidationerror, emitvalidationerror and responsevalidationerror errors.
118    """
119    if isinstance(e, RequestValidationError):
120        logger.error(f"request validation error: {e}")
121        return {'error': str(e)}
122    elif isinstance(e, ResponseValidationError):
123        logger.critical(f"response validation error: {e}")
124        raise e
125    if isinstance(e, EmitValidationError):
126        logger.critical(f"emit validation error: {e}")
127        raise e
128    else:
129        logger.critical(f"unknown error: {e}")
130        raise e
131
132if __name__ == '__main__':
133    # generate the asyncapi doc
134    path = pathlib.Path(__file__).parent / "chat_gpt_asyncapi.yaml"
135    doc_str = socketio.asyncapi_doc.get_yaml()
136    with open(path, "w") as f:
137        # doc_str = spec.get_json_str_doc()
138        f.write(doc_str)
139    # run the app
140    socketio.run(app, debug=True)
141

And here is the auto generated specification:

1asyncapi: 2.5.0
2channels:
3  /:
4    publish:
5      message:
6        oneOf:
7        - $ref: '#/components/messages/Create_Game'
8        - $ref: '#/components/messages/Make_Move'
9    subscribe:
10      message:
11        oneOf:
12        - $ref: '#/components/messages/game_created'
13        - $ref: '#/components/messages/move_made'
14        - $ref: '#/components/messages/game_drawn'
15        - $ref: '#/components/messages/game_won'
16    x-handlers:
17      disconnect: disconnect
18components:
19  messages:
20    Create_Game:
21      description: ''
22      name: create_game
23    Make_Move:
24      description: ''
25      name: make_move
26      payload:
27        $ref: '#/components/schemas/MakeMoveData'
28        deprecated: false
29      x-ack:
30        properties:
31          error:
32            description: The error message
33            example: Invalid move
34            title: Error
35            type: string
36        title: MakeMoveAckData
37        type: object
38    game_created:
39      description: ''
40      name: game_created
41      payload:
42        $ref: '#/components/schemas/GameCreatedData'
43        deprecated: false
44    game_drawn:
45      description: ''
46      name: game_drawn
47      payload:
48        $ref: '#/components/schemas/GameDrawnData'
49        deprecated: false
50    game_won:
51      description: ''
52      name: game_won
53      payload:
54        $ref: '#/components/schemas/GameWonData'
55        deprecated: false
56    move_made:
57      description: ''
58      name: move_made
59      payload:
60        $ref: '#/components/schemas/MoveMadeData'
61        deprecated: false
62  schemas:
63    GameCreatedData:
64      properties:
65        game_id:
66          description: The game id
67          example: 1
68          title: Game Id
69          type: integer
70      required:
71      - game_id
72      title: GameCreatedData
73      type: object
74    GameDrawnData:
75      properties: {}
76      title: GameDrawnData
77      type: object
78    GameWonData:
79      properties:
80        winner:
81          description: The winner
82          example: X
83          title: Winner
84          type: string
85      required:
86      - winner
87      title: GameWonData
88      type: object
89    MakeMoveAckData:
90      properties:
91        error:
92          description: The error message
93          example: Invalid move
94          title: Error
95          type: string
96      title: MakeMoveAckData
97      type: object
98    MakeMoveData:
99      properties:
100        game_id:
101          description: The game id
102          example: 1
103          title: Game Id
104          type: integer
105        x:
106          description: The x coordinate
107          example: 0
108          title: X
109          type: integer
110        y:
111          description: The y coordinate
112          example: 0
113          title: Y
114          type: integer
115      required:
116      - game_id
117      - x
118      - y
119      title: MakeMoveData
120      type: object
121    MoveMadeData:
122      properties:
123        board:
124          description: The game board
125          example:
126          - - X
127            - O
128            - ''
129          - - ''
130            - X
131            - ''
132          - - ''
133            - ''
134            - O
135          items:
136            items:
137              type: string
138            type: array
139          title: Board
140          type: array
141        turn:
142          description: The current turn
143          example: X
144          title: Turn
145          type: string
146      required:
147      - board
148      - turn
149      title: MoveMadeData
150      type: object
151    NoSpec:
152      deprecated: false
153      description: Specification is not provided
154info:
155  description: 'Tic-Tac-Toe Game API
156
157    <br/> AsyncAPI currently does not support Socket.IO binding and Web Socket like
158    syntax used for now.
159
160    In order to add support for Socket.IO ACK value, AsyncAPI is extended with with
161    x-ack keyword.
162
163    This documentation should **NOT** be used for generating code due to these limitations.
164
165    '
166  title: Tic-Tac-Toe API
167  version: 1.0.0
168servers:
169  TIC_TAC_TOE_BACKEND:
170    protocol: socketio
171    url: http://localhost:5000

Cover image by Windmills During Dawn from Unsplash

On this page