如何在 FastAPI 响应模式中使用箭头类型？

Question

I want to use Arrow type in FastAPI response because I am using it already in SQLAlchemy model (thanks to sqlalchemy_utils ).

I prepared a small self-contained example with a minimal FastAPI app. I expect that this app return product1 data from database.

Unfortunately the code below gives exception:

Exception has occurred: FastAPIError
Invalid args for response field! Hint: check that <class 'arrow.arrow.Arrow'> is a valid pydantic field type

import sqlalchemy
import uvicorn
from arrow import Arrow
from fastapi import FastAPI
from pydantic import BaseModel
from sqlalchemy import Column, Integer, Text, func
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy_utils import ArrowType

app = FastAPI()

engine = sqlalchemy.create_engine('sqlite:///db.db')
Base = declarative_base()

class Product(Base):
    __tablename__ = "product"
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(Text, nullable=True)
    created_at = Column(ArrowType(timezone=True), nullable=False, server_default=func.now())

Base.metadata.create_all(engine)


Session = sessionmaker(bind=engine)
session = Session()

product1 = Product(name="ice cream")
product2 = Product(name="donut")
product3 = Product(name="apple pie")

session.add_all([product1, product2, product3])
session.commit()


class ProductResponse(BaseModel):
    id: int
    name: str
    created_at: Arrow

    class Config:
        orm_mode = True
        arbitrary_types_allowed = True


@app.get('/', response_model=ProductResponse)
async def return_product():

    product = session.query(Product).filter(Product.id == 1).first()

    return product

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

requirements.txt:

sqlalchemy==1.4.23
sqlalchemy_utils==0.37.8
arrow==1.1.1
fastapi==0.68.1
uvicorn==0.15.0

This error is already discussed in those FastAPI issues:

1. https://github.com/tiangolo/fastapi/issues/1186
2. https://github.com/tiangolo/fastapi/issues/2382

---

One possible workaround is to add this code ( source ):

from pydantic import BaseConfig
BaseConfig.arbitrary_types_allowed = True

It is enough to put it just above @app.get('/'... , but it can be put even before app = FastAPI()

The problem with this solution is that output of GET endpoint will be:

// 20210826001330
// http://localhost:8000/

{
  "id": 1,
  "name": "ice cream",
  "created_at": {
    "_datetime": "2021-08-25T21:38:01+00:00"
  }
}

instead of desired:

// 20210826001330
// http://localhost:8000/

{
  "id": 1,
  "name": "ice cream",
  "created_at": "2021-08-25T21:38:01+00:00"
}

Answer 1

Add a custom function with the @validator decorator that returns the desired _datetime of the object:

class ProductResponse(BaseModel):
    id: int
    name: str
    created_at: Arrow

    class Config:
        orm_mode = True
        arbitrary_types_allowed = True

    @validator("created_at")
    def format_datetime(cls, value):
        return value._datetime

Tested on local, seems to be working:

$ curl -s localhost:8000 | jq
{
  "id": 1,
  "name": "ice cream",
  "created_at": "2021-12-02T08:25:10+00:00"
}

Answer 2

The solution is to monkeypatch pydantic's ENCODERS_BY_TYPE so it knows how to convert Arrow object so it can be accepted by json format:

from arrow import Arrow
from pydantic.json import ENCODERS_BY_TYPE
ENCODERS_BY_TYPE |= {Arrow: str}

Setting BaseConfig.arbitrary_types_allowed = True is also necessary.

Result:

// 20220514022717
// http://localhost:8000/

{
  "id": 1,
  "name": "ice cream",
  "created_at": "2022-05-14T00:20:11+00:00"
}

Full code:

import sqlalchemy
import uvicorn
from arrow import Arrow
from fastapi import FastAPI
from pydantic import BaseModel
from sqlalchemy import Column, Integer, Text, func
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy_utils import ArrowType

from pydantic.json import ENCODERS_BY_TYPE
ENCODERS_BY_TYPE |= {Arrow: str}

from pydantic import BaseConfig
BaseConfig.arbitrary_types_allowed = True

app = FastAPI()

engine = sqlalchemy.create_engine('sqlite:///db.db')
Base = declarative_base()

class Product(Base):
    __tablename__ = "product"
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(Text, nullable=True)
    created_at = Column(ArrowType(timezone=True), nullable=False, server_default=func.now())

Base.metadata.create_all(engine)


Session = sessionmaker(bind=engine)
session = Session()

product1 = Product(name="ice cream")
product2 = Product(name="donut")
product3 = Product(name="apple pie")

session.add_all([product1, product2, product3])
session.commit()


class ProductResponse(BaseModel):
    id: int
    name: str
    created_at: Arrow

    class Config:
        orm_mode = True
        arbitrary_types_allowed = True


@app.get('/', response_model=ProductResponse)
async def return_product():

    product = session.query(Product).filter(Product.id == 1).first()

    return product

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

Answer 3

Here is a code example where you do not need class Config and can work for any type by creating your own subclass with validators:

from psycopg2.extras import DateTimeTZRange as DateTimeTZRangeBase
from sqlalchemy.dialects.postgresql import TSTZRANGE
from sqlmodel import (
    Column,
    Field,
    Identity,
    SQLModel,
)

from pydantic.json import ENCODERS_BY_TYPE

ENCODERS_BY_TYPE |= {DateTimeTZRangeBase: str}


class DateTimeTZRange(DateTimeTZRangeBase):
    @classmethod
    def __get_validators__(cls):
        yield cls.validate

    @classmethod
    def validate(cls, v):
        if isinstance(v, str):
            lower = v.split(", ")[0][1:].strip().strip()
            upper = v.split(", ")[1][:-1].strip().strip()
            bounds = v[:1] + v[-1:]
            return DateTimeTZRange(lower, upper, bounds)
        elif isinstance(v, DateTimeTZRangeBase):
            return v
        raise TypeError("Type must be string or DateTimeTZRange")

    @classmethod
    def __modify_schema__(cls, field_schema):
        field_schema.update(type="string", example="[2022,01,01, 2022,02,02)")


class EventBase(SQLModel):
    __tablename__ = "event"
    timestamp_range: DateTimeTZRange = Field(
        sa_column=Column(
            TSTZRANGE(),
            nullable=False,
        ),
    )


class Event(EventBase, table=True):
    id: int | None = Field(
        default=None,
        sa_column_args=(Identity(always=True),),
        primary_key=True,
        nullable=False,
    )

link to Github issue: https://github.com/tiangolo/sqlmodel/issues/235#issuecomment-1162063590