python使用泛型

所谓的泛型, 就是将数据类型作为参数进行传递, 即在我们用的时候确定数据类型, 这是一种在面向对象语言中经常使用的特性

一般类使用#

以SQLAlchemy举例

比如: 我们统一写个将数据保存到数据库的接口, 只有将数据库链接 表对象 数据传入即可, 返回的是表对象的实例, 为了让IDE可以识别返回对象, 我们可以使用泛型
这里需要用到:

1. typing的TypeVar和Type
TypeVar是类型变量, 主要用于泛型类型与泛型函数定义的参数, 第一个参数是名称, bound参数用于规定该类型为bound值的子类
   Type[C]的形式为协变量, 表明C的所有子类都应 使用与C相同的 构造器签名 及 类方法签名, 假如我们需要返回泛型类型的话, 需要用到

   更多使用方法, 见: typing使用

2. 使用了pydantic规范要创建数据的类型

   关于pydantic的一般使用, 见: Pydantic使用

from typing import TypeVar, Type
from sqlalchemy.orm import Session
from pydantic import BaseModel
from orm.models import Category
from orm.schemas import WriteCategoryModel
 
# 定义类型
ModelT = TypeVar("ModelT")
DataT = TypeVar("DataT", bound=BaseModel)  # bound表明该类为BaseModel的子类
 
"""
为节省空间, 表结构和模型不展示
"""
def create_category(session: Session, data: WriteCategoryModel) -> Type[Category]:
    cate_obj = save_one_to_db(session=session, model_class=Category, data=data)
    return cate_obj
 
 
def save_one_to_db(session: Session, model_class: ModelT, data: DataT) -> ModelT:
    """
    保存一条到数据库
    :param session: SQLAlchemy Session
    :param model_class: sqlalchemy模型类
    :param data: pydantic模型对象
    :return: 对应sqlalchemy模型类的对象
    """
    try:
        obj = model_class(**data.dict())
        session.add(obj)
        session.commit()
        # 手动将 数据 刷新到数据库
        session.refresh(obj)
        return obj
 
    except Exception as e:
        # 别忘记发生错误时回滚
        session.rollback()
        raise e
 

pydantic使用#

在使用pydantic时, 亦可以使用泛型

比如: 在FastAPI中返回统一返回格式为:

{
   "status": true,
   "msg": "success",
   "data": ...
}

我们的data可能是列表或对象, 而且对应的pydantic模型也不一样, 这时我们就可以使用泛型了

代码:

from typing import List, Optional, Generic, TypeVar
 
from fastapi import APIRouter, status, HTTPException
from pydantic import BaseModel
from pydantic.generics import GenericModel
 
router = APIRouter(prefix="/test", tags=["测试泛型"])
 
# 为了方便, 在这里定义pydantic模型
DataT = TypeVar("DataT")
 
 
class GenericResponse(GenericModel, Generic[DataT]):
    """
    通用返回数据
    """
    status: bool
    msg: str
    data: Optional[DataT] = None  # 可能连data都没有
    # 设置response_model_exclude_unset=True即可
 
 
class BookModel(BaseModel):
    id: int
    name: str
 
 
# 伪数据
fake_book_data = [
    {"id": 1, "name": "book1"},
    {"id": 2, "name": "book2"},
    {"id": 3, "name": "book3"},
    {"id": 4, "name": "book4"},
    {"id": 5, "name": "book5"},
]
 
 
@router.get("/books", response_model=GenericResponse[List[BookModel]])
def get_books():
    return {
        "status": True,
        "msg": "获取成功",
        "data": fake_book_data
    }
 
 
@router.get("/books/{bid}", response_model=GenericResponse[BookModel])
def retrieve_book(bid: int):
    for item in fake_book_data:
        if item.get("id") == bid:
            return {
                "status": True,
                "msg": "获取成功",
                "data": item
            }
    # 不存在
    raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="该书不存在")
 

访问/docs页面, 成功通过测试