如何为msgspec Struct属性添加注释到JSON Schema或改用Pydantic？

你期望的JSON Schema格式是合规的，description是JSON Schema标准支持的属性，用于为结构或字段添加说明文本。以下是分别基于msgspec和Pydantic的实现方案：

一、msgspec 实现方案 #

msgspec默认不支持提取属性的行内注释，需要通过AST解析类定义的源码，手动提取注释后修改生成的JSON Schema：

import json
import ast
import msgspec
from msgspec import Struct

def print_schema(schema):
    encoded_schema = msgspec.json.encode(schema)
    formatted_schema = json.dumps(json.loads(encoded_schema), indent=4)
    print(formatted_schema)

def get_field_comments(cls):
    """提取Struct类中属性的行内注释"""
    tree = ast.parse(open(cls.__module__.__file__).read())
    class_node = None
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef) and node.name == cls.__name__:
            class_node = node
            break
    if not class_node:
        return {}
    
    comments = {}
    for stmt in class_node.body:
        if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name):
            field_name = stmt.target.id
            if stmt.end_lineno:
                with open(cls.__module__.__file__, 'r') as f:
                    lines = f.readlines()
                comment_line = lines[stmt.end_lineno - 1].strip()
                if comment_line.startswith('#'):
                    comments[field_name] = comment_line[1:].strip()
    return comments

class Product(Struct):
    """A product in a catalog"""
    id: int      # DB uid
    name: str    # Name of product
    price: float # Price of product

# 生成原始Schema并注入注释
schema = msgspec.json.schema(Product)
field_comments = get_field_comments(Product)
product_def = schema['$defs']['Product']
for field, desc in field_comments.items():
    if field in product_def['properties']:
        product_def['properties'][field]['description'] = desc

print_schema(schema)

注意：该方法依赖本地源码文件，无法在交互式环境（如Jupyter Notebook）中使用。

二、Pydantic 实现方案 #

Pydantic原生支持字段描述，有两种常用实现方式：

方式1：使用Field显式指定描述 #

这是最可靠、最符合Pydantic规范的写法：

import json
from pydantic import BaseModel, Field

class Product(BaseModel):
    """A product in a catalog"""
    id: int = Field(description="DB uid")
    name: str = Field(description="Name of product")
    price: float = Field(description="Price of product")

# 生成并打印JSON Schema
schema = Product.model_json_schema()
print(json.dumps(schema, indent=4))

方式2：自动提取行内注释作为描述 #

如果希望保持代码简洁，可通过AST解析自动提取注释：

import json
import ast
from pydantic import BaseModel

def set_field_descriptions(cls):
    """为BaseModel字段添加行内注释作为description"""
    tree = ast.parse(open(cls.__module__.__file__).read())
    class_node = None
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef) and node.name == cls.__name__:
            class_node = node
            break
    if not class_node:
        return cls
    
    comments = {}
    for stmt in class_node.body:
        if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name):
            field_name = stmt.target.id
            if stmt.end_lineno:
                with open(cls.__module__.__file__, 'r') as f:
                    lines = f.readlines()
                comment_line = lines[stmt.end_lineno - 1].strip()
                if comment_line.startswith('#'):
                    comments[field_name] = comment_line[1:].strip()
    
    for field_name, desc in comments.items():
        if field_name in cls.model_fields:
            cls.model_fields[field_name].description = desc
    return cls

@set_field_descriptions
class Product(BaseModel):
    """A product in a catalog"""
    id: int      # DB uid
    name: str    # Name of product
    price: float # Price of product

schema = Product.model_json_schema()
print(json.dumps(schema, indent=4))

内容的提问来源于stack exchange，提问作者user2746401