feat: ✨ functions/generators as handler dependencies (#35)

pest/core/application.py

@@ -237,6 +237,9 @@ def enable_cors(self, **opts: Unpack[CorsOptions]) -> None:
 
         self.add_middleware(CORSMiddleware, **opts)
 
+    def add_middleware(self, middleware_class: type, *args: Any, **kwargs: Any) -> None:
+        super().add_middleware(middleware_class, *args, **kwargs)
+
     def build_middleware_stack(self) -> ASGIApp:
         # Duplicate/override from FastAPI to add the di_scope_middleware, which
         # is required for Pest's DI to work. We need it to run before any other

pest/core/handler.py

@@ -1,5 +1,5 @@
 from dataclasses import asdict
-from inspect import Parameter, signature
+from inspect import Parameter, isclass, isfunction, signature
 from typing import TYPE_CHECKING, Any, Callable, List, Tuple, Type, Union, get_args
 
 try:
@@ -98,11 +98,22 @@ def _get_new_param(ctrl: type, parameter: Parameter) -> Parameter:
 
             annotation = pest_anns[0]
             if annotation.token is not None:
-                # we replace the parameter with a `Depends` object that resolves the value
-                # through the `module`'s container
-                parameter = parameter.replace(
-                    default=Depends(PestFastAPIInjector(controller=ctrl, token=annotation.token))
-                )
+                # if the token is a function, we replace it with a FastAPI's `Depends(dep)` call
+                if isfunction(annotation.token):
+                    parameter = parameter.replace(default=Depends(annotation.token))
+                elif isclass(annotation.token):
+                    # otherwise we replace the parameter with a `Depends` on `PestFastAPIInjector`
+                    # which will try to resolve the value from the `module`'s container
+                    parameter = parameter.replace(
+                        default=Depends(
+                            PestFastAPIInjector(controller=ctrl, token=annotation.token)
+                        )
+                    )
+                else:
+                    raise PestException(
+                        'Invalid injection annotation token!',
+                        hint=f'Parameter {parameter.name} has an invalid injection token!',
+                    )
     return parameter.replace(kind=Parameter.KEYWORD_ONLY)
 
 

pest/di/injection.py

@@ -1,4 +1,4 @@
-from typing import Type, TypeVar, Union, cast
+from typing import Callable, Type, TypeVar, Union, cast
 
 T = TypeVar('T')
 
@@ -9,11 +9,11 @@ class _Inject:
     @internal you should not use this class directly, use `inject` instead.
     """
 
-    def __init__(self, token: Union[type, None]) -> None:
+    def __init__(self, token: Union[Type, None, Callable]) -> None:
         self.token = token
 
 
-def inject(token: Union[None, Type[T]] = None) -> T:
+def inject(token: Union[None, Type[T], Callable] = None) -> T:
     """
     🐀 ⇝ marks a parameter as a dependency to be injected by `pest` 💉
 

tests/cfg/test_modules/rodi_route_dependencies_functions.py

@@ -0,0 +1,53 @@
+try:
+    from typing import Annotated
+except ImportError:
+    from typing_extensions import Annotated
+
+
+from typing import Union
+from uuid import uuid4
+
+from pest.decorators.controller import controller
+from pest.decorators.handler import get
+from pest.decorators.module import module
+from pest.di import injected
+
+
+def who_am_i(user: Union[str, None] = None):
+    return user if user else str(uuid4())
+
+
+def yield_me(user: Union[str, None] = None):
+    yield user if user else str(uuid4())
+
+
+@controller('')
+class FooController:
+    @get('/assigned')
+    def assigned(self, me: str = injected(who_am_i)):
+        return {'id': me}
+
+    @get('/assigned-with-yield')
+    def assign_with_yield(self, me: str = injected(yield_me)):
+        return {'id': me}
+
+
+@controller('')
+class FooController39plus:
+    @get('/assigned')
+    def assigned(self, me: Annotated[str, injected(who_am_i)]):
+        return {'id': me}
+
+    @get('/assigned-with-yield')
+    def assign_with_yield(self, me: Annotated[str, injected(yield_me)]):
+        return {'id': me}
+
+
+@module(controllers=[FooController])
+class FunctionsDependenciesModule:
+    pass
+
+
+@module(controllers=[FooController39plus])
+class FunctionsDependenciesModule39plus:
+    pass

tests/test_handler.py

@@ -12,6 +12,10 @@
     RodiDependenciesModule,
     RodiDependenciesModule39plus,
 )
+from .cfg.test_modules.rodi_route_dependencies_functions import (
+    FunctionsDependenciesModule,
+    FunctionsDependenciesModule39plus,
+)
 
 
 def test_get_handler():
@@ -141,7 +145,7 @@ def test_handler_can_inject_di() -> None:
 
 @pytest.mark.skipif(sys.version_info < (3, 9), reason='requires python3.9 or higher')
 def test_handlder_can_inject_di_annotation() -> None:
-    """🐀 handlers :: di :: should be able to be injected by rodi"""
+    """🐀 handlers :: di :: should be able to be injected by rodi using Annotation"""
 
     app = Pest.create(RodiDependenciesModule39plus)
     client = TestClient(app)
@@ -150,3 +154,65 @@ def test_handlder_can_inject_di_annotation() -> None:
     assert response.status_code == 200
     assert isinstance(response.json().get('id'), str)
     assert len(response.json().get('id')) > 0
+
+
+def test_handler_can_inject_functional_deps() -> None:
+    """🐀 handlers :: di :: should be able to have functions and generators as dependencies"""
+    app = Pest.create(FunctionsDependenciesModule)
+    client = TestClient(app)
+
+    response = client.get('/assigned')
+    assert response.status_code == 200
+    assert isinstance(response.json().get('id'), str)
+    assert len(response.json().get('id')) > 0
+
+    # function dependencies should also be able to get query params
+    response = client.get('/assigned?user=foo')
+    assert response.status_code == 200
+    assert isinstance(response.json().get('id'), str)
+    assert response.json().get('id') == 'foo'
+
+    # should also work with generators
+    response = client.get('/assigned-with-yield')
+    assert response.status_code == 200
+    assert isinstance(response.json().get('id'), str)
+    assert len(response.json().get('id')) > 0
+
+    # generator function dependencies should also be able to get query params
+    response = client.get('/assigned-with-yield?user=foo')
+    assert response.status_code == 200
+    assert isinstance(response.json().get('id'), str)
+    assert response.json().get('id') == 'foo'
+
+
+@pytest.mark.skipif(sys.version_info < (3, 9), reason='requires python3.9 or higher')
+def test_handler_can_inject_functional_deps_annotations() -> None:
+    """
+    🐀 handlers :: di :: should be able to have functions and generators as annotated
+                         dependencies
+    """
+    app = Pest.create(FunctionsDependenciesModule39plus)
+    client = TestClient(app)
+
+    response = client.get('/assigned')
+    assert response.status_code == 200
+    assert isinstance(response.json().get('id'), str)
+    assert len(response.json().get('id')) > 0
+
+    # function dependencies should also be able to get query params
+    response = client.get('/assigned?user=foo')
+    assert response.status_code == 200
+    assert isinstance(response.json().get('id'), str)
+    assert response.json().get('id') == 'foo'
+
+    # should also work with generators
+    response = client.get('/assigned-with-yield')
+    assert response.status_code == 200
+    assert isinstance(response.json().get('id'), str)
+    assert len(response.json().get('id')) > 0
+
+    # generator function dependencies should also be able to get query params
+    response = client.get('/assigned-with-yield?user=foo')
+    assert response.status_code == 200
+    assert isinstance(response.json().get('id'), str)
+    assert response.json().get('id') == 'foo'