Support `dataclasses` in responses (#3576)

Co-authored-by: amit lissack <amit@opentrons.com>
This commit is contained in:
Sebastián Ramírez 2021-07-21 16:39:12 +02:00 committed by GitHub
parent 49517d3310
commit 96fdfc53cc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 66 additions and 0 deletions

View File

@ -1,3 +1,4 @@
import dataclasses
from collections import defaultdict
from enum import Enum
from pathlib import PurePath
@ -61,6 +62,8 @@ def jsonable_encoder(
custom_encoder=encoder,
sqlalchemy_safe=sqlalchemy_safe,
)
if dataclasses.is_dataclass(obj):
return dataclasses.asdict(obj)
if isinstance(obj, Enum):
return obj.value
if isinstance(obj, PurePath):

View File

@ -1,4 +1,5 @@
import asyncio
import dataclasses
import email.message
import enum
import inspect
@ -90,6 +91,8 @@ def _prepare_response_content(
)
for k, v in res.items()
}
elif dataclasses.is_dataclass(res):
return dataclasses.asdict(res)
return res

View File

@ -19,6 +19,11 @@ def get_valid():
return {"name": "valid", "price": 1.0}
@app.get("/items/object", response_model=Item)
def get_object():
return Item(name="object", price=1.0, owner_ids=[1, 2, 3])
@app.get("/items/coerce", response_model=Item)
def get_coerce():
return {"name": "coerce", "price": "1.0"}
@ -33,6 +38,29 @@ def get_validlist():
]
@app.get("/items/objectlist", response_model=List[Item])
def get_objectlist():
return [
Item(name="foo"),
Item(name="bar", price=1.0),
Item(name="baz", price=2.0, owner_ids=[1, 2, 3]),
]
@app.get("/items/no-response-model/object")
def get_no_response_model_object():
return Item(name="object", price=1.0, owner_ids=[1, 2, 3])
@app.get("/items/no-response-model/objectlist")
def get_no_response_model_objectlist():
return [
Item(name="foo"),
Item(name="bar", price=1.0),
Item(name="baz", price=2.0, owner_ids=[1, 2, 3]),
]
client = TestClient(app)
@ -42,6 +70,12 @@ def test_valid():
assert response.json() == {"name": "valid", "price": 1.0, "owner_ids": None}
def test_object():
response = client.get("/items/object")
response.raise_for_status()
assert response.json() == {"name": "object", "price": 1.0, "owner_ids": [1, 2, 3]}
def test_coerce():
response = client.get("/items/coerce")
response.raise_for_status()
@ -56,3 +90,29 @@ def test_validlist():
{"name": "bar", "price": 1.0, "owner_ids": None},
{"name": "baz", "price": 2.0, "owner_ids": [1, 2, 3]},
]
def test_objectlist():
response = client.get("/items/objectlist")
response.raise_for_status()
assert response.json() == [
{"name": "foo", "price": None, "owner_ids": None},
{"name": "bar", "price": 1.0, "owner_ids": None},
{"name": "baz", "price": 2.0, "owner_ids": [1, 2, 3]},
]
def test_no_response_model_object():
response = client.get("/items/no-response-model/object")
response.raise_for_status()
assert response.json() == {"name": "object", "price": 1.0, "owner_ids": [1, 2, 3]}
def test_no_response_model_objectlist():
response = client.get("/items/no-response-model/objectlist")
response.raise_for_status()
assert response.json() == [
{"name": "foo", "price": None, "owner_ids": None},
{"name": "bar", "price": 1.0, "owner_ids": None},
{"name": "baz", "price": 2.0, "owner_ids": [1, 2, 3]},
]