Add @endpoint.response_model

apphelpers/rest/endpoint.py

@@ -46,3 +46,11 @@ def not_found_on_none(func):
 def ignore_site_ctx(func):
     func.ignore_site_ctx = True
     return func
+
+
+def response_model(response_model):
+    def decorator(func):
+        func.response_model = response_model
+        return func
+
+    return decorator

apphelpers/rest/fastapi.py

@@ -1,5 +1,6 @@
 import inspect
 from functools import wraps
+from types import NoneType, UnionType
 
 from converge import settings
 from fastapi import APIRouter, Depends
@@ -22,6 +23,12 @@
     from honeybadger.utils import filter_dict
 
 
+def _is_optional(t) -> bool:
+    return (
+        isinstance(t, UnionType) and len(t.__args__) == 2 and t.__args__[1] == NoneType
+    )
+
+
 def raise_not_found_on_none(f):
     if getattr(f, "not_found_on_none", None) is True:
 
@@ -470,17 +477,22 @@ def build(self, method, method_args, method_kw, f):
         module = f.__module__.split(".")[-1].strip("_")
         name = f.__name__.strip("_")
         return_type = inspect.signature(f).return_annotation
+        response_model = getattr(f, "response_model", return_type)
 
         if "operation_id" not in method_kw:
             method_kw["operation_id"] = f"{name}_{module}"
         if "tags" not in method_kw:
             method_kw["tags"] = [module]
+
         if (
             "response_model" not in method_kw
             and "response_class" not in method_kw
-            and return_type is not inspect.Signature.empty
+            and response_model is not inspect.Signature.empty
         ):
-            method_kw["response_model"] = return_type
+            method_kw["response_model"] = response_model
+
+        if _is_optional(return_type) and not _is_optional(response_model):
+            f.not_found_on_none = True
 
         print(
             f"{method_args[0]}",

fastapi_tests/app/endpoints.py

@@ -1,5 +1,5 @@
 from apphelpers.rest import endpoint as ep
-from apphelpers.rest.fastapi import json_body, user, user_id, user_agent
+from apphelpers.rest.fastapi import json_body, user, user_agent, user_id
 
 
 def echo(word, user=user):
@@ -46,6 +46,12 @@ def get_snake(name=None):
     return name
 
 
+@ep.login_required
+@ep.response_model(str)
+def get_snake_fancy(name=None) -> str | None:
+    return name
+
+
 @ep.login_required
 @ep.not_found_on_none
 async def get_snake_async(name=None):
@@ -91,6 +97,8 @@ def setup_routes(factory):
     factory.get("/snakes/{name}")(get_snake)
     factory.get("/snakes-async/{name}")(get_snake_async)
 
+    factory.get("/snakes-fancy/{name}")(get_snake_fancy)
+
     factory.get("/sites/{site_id}/echo-groups")(echo_site_groups)
     factory.get("/sites/{site_id}/echo-groups-async")(echo_site_groups_async)
 

fastapi_tests/test_rest.py

@@ -163,9 +163,15 @@ def test_not_found_on_none():
     url = base_url + "snakes/viper"
     assert requests.get(url).status_code != 404
 
+    url = base_url + "snakes-fancy/viper"
+    assert requests.get(url).status_code != 404
+
     url = base_url + "snakes"
     assert requests.get(url).status_code == 404
 
+    url = base_url + "snakes-fancy"
+    assert requests.get(url).status_code == 404
+
 
 def test_site_group_access():
     url = echo_site_groups_url