Add support for exceptions raised in PEP-593 Annotated dependencies - new FastAPI feature - fastapi/fastapi#4871

Saran33 · Saran33 · commit 250c32ac1e75 · 2023-04-02T22:35:23.000+01:00
diff --git a/fastapi_docx/exception_finder.py b/fastapi_docx/exception_finder.py
@@ -2,6 +2,7 @@
 import importlib
 import inspect
 import textwrap
+import typing
 from collections.abc import Callable, Generator, Iterable
 from itertools import chain
 from types import ModuleType
@@ -56,6 +57,21 @@ def is_callable_instance(obj: object) -> bool:
     return hasattr(obj, "__call__") and not isinstance(obj, type)
 
 
+def is_annotated_alias(annot):
+    _AnnotatedAlias = getattr(typing, "_AnnotatedAlias")
+    return isinstance(annot, _AnnotatedAlias)
+
+
+def get_annotated_dependency(annot):
+    if is_annotated_alias(annot):
+        annot_args = typing.get_args(annot)
+        annot_dep = annot_args[1] if len(annot_args) > 1 else None
+        if annot_dep and annot_dep.__class__.__name__ == "Depends":
+            dependency = getattr(annot_dep, "dependency", None)
+            return dependency
+    return None
+
+
 def create_exc_instance(
     exc_class: type[Exception], exc_args: list[Any] | None = None
 ) -> Exception | None:
@@ -104,6 +120,8 @@ def extract_exceptions(
             self.exceptions.extend(self.find_exceptions(function))
         if self.dependencyClasses:
             self.exceptions += self.find_dependency_exceptions(route)
+        if self.dependencyClasses:
+            self.exceptions += self.find_annotated_dependency_exceptions(route)
         if self.serviceClasses:
             self.exceptions += self.find_service_exceptions(route)
         return self.exceptions
@@ -241,6 +259,35 @@ def find_dependency_exceptions(
                                 exceptions.extend(_exceptions)
         return exceptions
 
+    def find_annotated_dependency_exceptions(
+        self,
+        route: APIRoute | Callable,
+    ) -> list[HTTPException]:
+        exceptions = []
+        assert self.dependencyClasses is not None
+        func = route.endpoint if hasattr(route, "endpoint") else route
+        module = inspect.getmodule(func)
+        source = inspect.getsource(func)
+        tree = ast.parse(source)
+
+        for node in find_nodes(tree, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            for kwarg in getattr(node.args, "kwonlyargs", []):
+                if kwarg and hasattr(kwarg, "annotation"):
+                    if (
+                        annot := kwarg.annotation.id
+                        if hasattr(kwarg.annotation, "id") and kwarg.annotation.id
+                        else None
+                    ):
+                        cls = None
+                        try:
+                            cls = getattr(module, annot)
+                        except (AttributeError, NameError):
+                            ...
+                        if cls and (dependency := get_annotated_dependency(cls)):
+                            if _exceptions := self.find_exceptions(dependency, module):
+                                exceptions.extend(_exceptions)
+        return exceptions
+
     def create_exc_inst_from_raise_stmt(
         self,
         raise_stmt: ast.Raise,
diff --git a/tests/unit_tests/test_dependency_exceptions.py b/tests/unit_tests/test_dependency_exceptions.py
@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Annotated, Any
 
 from fastapi import APIRouter, Depends, FastAPI
 
@@ -13,11 +13,11 @@
 
 router = APIRouter()
 
+UserDict = Annotated[dict[str, Any], Depends(UserDeps.get_current_user)]
+
 
 @router.get("/me")
-def get_user(
-    *, user_in, current_user: dict[str, Any] = Depends(UserDeps.get_current_user)
-):
+def get_user(*, user_in, current_user: UserDict):
     result = UserService.get_authenticated(user_in)
     return result
 
