如何为Python Flask与Marshmallow添加VSCode类型推断支持

首先注意：你的PostInput需要继承marshmallow.Schema才能成为有效的验证Schema，以下是几种让VSCode能正确推断类型的方案：

方法1：用TypedDict手动标注返回类型 #

先定义一个和Schema结构匹配的TypedDict，给load返回的data加上类型注解，VSCode就能识别字段类型：

from typing import TypedDict
from marshmallow import Schema, fields, ValidationError
from flask import Flask, request, jsonify

app = Flask(__name__)

class PostInputSchema(Schema):
    myString = fields.Str(required=True)

# 定义与Schema对应的类型结构
class PostInputData(TypedDict):
    myString: str

@app.route('/foo', methods=['POST'])
def post_foo():
    try:
        # 给data添加类型注解
        data: PostInputData = PostInputSchema().load(request.json)
    except ValidationError as err:
        return jsonify(err.messages), 400
      
    myString = data.get("myString")
    # 此时VSCode会识别myString为str类型
    return jsonify({"received": myString})

方法2：用marshmallow-dataclass（贴近TS的DTO体验） #

这个库可以让你用Python dataclass定义数据结构，自动生成Marshmallow Schema，同时保留完整类型信息，VSCode的类型推断直接生效：

先安装依赖：

pip install marshmallow-dataclass

代码示例：

from marshmallow_dataclass import dataclass
from marshmallow import ValidationError
from flask import Flask, request, jsonify

app = Flask(__name__)

# 用dataclass定义DTO，自动生成验证Schema
@dataclass
class PostInput:
    myString: str  # 类型注解直接作为验证规则，默认必填

@app.route('/foo', methods=['POST'])
def post_foo():
    try:
        # load返回PostInput实例，类型完全可推断
        data = PostInput.Schema().load(request.json)
    except ValidationError as err:
        return jsonify(err.messages), 400
      
    # 直接访问实例属性，VSCode能识别myString为str类型
    myString = data.myString
    return jsonify({"received": myString})

这种方式和TypeScript定义interface的体验最接近，不用重复编写Schema和类型定义。

方法3：给Schema添加Generic类型支持 #

如果不想用额外库，可以借助Marshmallow的Generic扩展来指定返回类型：

from typing import Generic, TypeVar, TypedDict
from marshmallow import Schema, fields, ValidationError
from flask import Flask, request, jsonify

app = Flask(__name__)

T = TypeVar('T')

class TypedSchema(Schema, Generic[T]):
    def load(self, data, **kwargs) -> T:
        return super().load(data, **kwargs)

class PostInputData(TypedDict):
    myString: str

class PostInputSchema(TypedSchema[PostInputData]):
    myString = fields.Str(required=True)

@app.route('/foo', methods=['POST'])
def post_foo():
    try:
        data = PostInputSchema().load(request.json)
    except ValidationError as err:
        return jsonify(err.messages), 400
      
    myString = data.get("myString")
    # VSCode会识别myString为str类型
    return jsonify({"received": myString})

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