add response models to each return type

setup.cfg

@@ -29,6 +29,7 @@ install_requires =
     uvicorn
     pymongo
     motor
+    setuptools==63.4.2
 
 [options.extras_require]
 # For development tests/docs

src/diffcalc_API/models/__init__.py

@@ -1,3 +1,3 @@
-from diffcalc_API.models import hkl, ub
+from diffcalc_API.models import hkl, response, ub
 
-__all__ = ["ub", "hkl"]
+__all__ = ["ub", "hkl", "response"]

src/diffcalc_API/models/response.py

@@ -0,0 +1,29 @@
+from typing import Dict, List
+
+from pydantic import BaseModel
+
+from diffcalc_API.models.ub import HklModel
+
+
+class InfoResponse(BaseModel):
+    message: str
+
+
+class StringResponse(BaseModel):
+    payload: str
+
+
+class ArrayResponse(BaseModel):
+    payload: List[List[float]]
+
+
+class ScanResponse(BaseModel):
+    payload: Dict[str, List[Dict[str, float]]]
+
+
+class DiffractorAnglesResponse(BaseModel):
+    payload: List[Dict[str, float]]
+
+
+class MillerIndicesResponse(BaseModel):
+    payload: HklModel

src/diffcalc_API/routes/constraints.py

@@ -2,40 +2,40 @@
 
 from fastapi import APIRouter, Body, Depends, Query
 
+from diffcalc_API.models.response import InfoResponse, StringResponse
 from diffcalc_API.services import constraints as service
 from diffcalc_API.stores.protocol import HklCalcStore, get_store
 
 router = APIRouter(prefix="/constraints", tags=["constraints"])
 
 
-@router.get("/{name}")
+@router.get("/{name}", response_model=StringResponse)
 async def get_constraints(
     name: str,
     store: HklCalcStore = Depends(get_store),
     collection: Optional[str] = Query(default=None, example="B07"),
 ):
     content = await service.get_constraints(name, store, collection)
-    return {"payload": content}
+    return StringResponse(payload=content)
 
 
-@router.post("/{name}")
+@router.post("/{name}", response_model=InfoResponse)
 async def set_constraints(
     name: str,
     constraints: Dict[str, float] = Body(example={"qaz": 0, "alpha": 0, "eta": 0}),
     store: HklCalcStore = Depends(get_store),
     collection: Optional[str] = Query(default=None, example="B07"),
 ):
     await service.set_constraints(name, constraints, store, collection)
-
-    return {
-        "message": (
+    return InfoResponse(
+        message=(
             f"constraints updated (replaced) for crystal {name} in "
             + f"collection {collection}"
         )
-    }
+    )
 
 
-@router.delete("/{name}/{property}")
+@router.delete("/{name}/{property}", response_model=InfoResponse)
 async def remove_constraint(
     name: str,
     property: str,
@@ -44,15 +44,15 @@ async def remove_constraint(
 ):
     await service.remove_constraint(name, property, store, collection)
 
-    return {
-        "message": (
+    return InfoResponse(
+        message=(
             f"unconstrained {property} for crystal {name} in "
             + f"collection {collection}. "
         )
-    }
+    )
 
 
-@router.patch("/{name}/{property}")
+@router.patch("/{name}/{property}", response_model=InfoResponse)
 async def set_constraint(
     name: str,
     property: str,
@@ -62,9 +62,9 @@ async def set_constraint(
 ):
     await service.set_constraint(name, property, value, store, collection)
 
-    return {
-        "message": (
+    return InfoResponse(
+        message=(
             f"constrained {property} for crystal {name} in collection "
             + f"{collection}. "
         )
-    }
+    )

src/diffcalc_API/routes/hkl.py

@@ -4,14 +4,20 @@
 
 from diffcalc_API.errors.hkl import InvalidSolutionBoundsError
 from diffcalc_API.models.hkl import SolutionConstraints
+from diffcalc_API.models.response import (
+    ArrayResponse,
+    DiffractorAnglesResponse,
+    MillerIndicesResponse,
+    ScanResponse,
+)
 from diffcalc_API.models.ub import HklModel, PositionModel
 from diffcalc_API.services import hkl as service
 from diffcalc_API.stores.protocol import HklCalcStore, get_store
 
 router = APIRouter(prefix="/hkl", tags=["hkl"])
 
 
-@router.get("/{name}/UB")
+@router.get("/{name}/UB", response_model=ArrayResponse)
 async def calculate_ub(
     name: str,
     tag1: Optional[str] = Query(default=None, example="refl1"),
@@ -24,10 +30,10 @@ async def calculate_ub(
     content = await service.calculate_ub(
         name, store, collection, tag1, idx1, tag2, idx2
     )
-    return {"payload": content}
+    return ArrayResponse(payload=content)
 
 
-@router.get("/{name}/position/lab")
+@router.get("/{name}/position/lab", response_model=DiffractorAnglesResponse)
 async def lab_position_from_miller_indices(
     name: str,
     miller_indices: HklModel = Depends(),
@@ -50,11 +56,10 @@ async def lab_position_from_miller_indices(
         store,
         collection,
     )
+    return DiffractorAnglesResponse(payload=positions)
 
-    return {"payload": positions}
 
-
-@router.get("/{name}/position/hkl")
+@router.get("/{name}/position/hkl", response_model=MillerIndicesResponse)
 async def miller_indices_from_lab_position(
     name: str,
     pos: PositionModel = Depends(),
@@ -65,10 +70,10 @@ async def miller_indices_from_lab_position(
     hkl = await service.miller_indices_from_lab_position(
         name, pos, wavelength, store, collection
     )
-    return {"payload": hkl}
+    return MillerIndicesResponse(payload=hkl)
 
 
-@router.get("/{name}/scan/hkl")
+@router.get("/{name}/scan/hkl", response_model=ScanResponse)
 async def scan_hkl(
     name: str,
     start: List[float] = Query(..., example=[1, 0, 1]),
@@ -95,10 +100,10 @@ async def scan_hkl(
         store,
         collection,
     )
-    return {"payload": scan_results}
+    return ScanResponse(payload=scan_results)
 
 
-@router.get("/{name}/scan/wavelength")
+@router.get("/{name}/scan/wavelength", response_model=ScanResponse)
 async def scan_wavelength(
     name: str,
     start: float = Query(..., example=1.0),
@@ -118,10 +123,10 @@ async def scan_wavelength(
     scan_results = await service.scan_wavelength(
         name, start, stop, inc, hkl, solution_constraints, store, collection
     )
-    return {"payload": scan_results}
+    return ScanResponse(payload=scan_results)
 
 
-@router.get("/{name}/scan/{constraint}")
+@router.get("/{name}/scan/{constraint}", response_model=ScanResponse)
 async def scan_constraint(
     name: str,
     constraint: str,
@@ -153,4 +158,4 @@ async def scan_constraint(
         collection,
     )
 
-    return {"payload": scan_results}
+    return ScanResponse(payload=scan_results)