Added LRU cache to WireupTask to prevent recompilation (#127)

* Added LRU cache to WireupTask to prevent recompilation

* Fixed ruff formatting

* Added the recommended changes

* Added the recommended changes

* Updated the recommended changes in the .md file and also added tests

* Fixed error; moved helper functions to module level for caching

* Made the ruff formatting changes

Author: sap6011

docs/pages/class/fastapi_integration.md

@@ -1,3 +1,8 @@
 ## wireup.integration.fastapi
 
 ::: wireup.integration.fastapi
+
+!!! note
+    For best performance with `WireupTask`, prefer regular top-level functions for background task callbacks.
+    Wireup creates cached injection wrappers for top-level functions, but callable objects and nested
+    functions/closures do not benefit from that caching.

docs/pages/class/starlette_integration.md

@@ -1,3 +1,8 @@
 ## wireup.integration.flask
 
 ::: wireup.integration.starlette
+
+!!! note
+    For best performance with `WireupTask`, prefer regular top-level functions for background task callbacks.
+    Wireup creates cached injection wrappers for top-level functions, but callable objects and nested
+    functions/closures do not benefit from that caching.

test/integration/starlette/test_starlette_integration.py

@@ -1,3 +1,4 @@
+import types
 from typing import Iterator
 from uuid import uuid4
 
@@ -278,3 +279,68 @@ def test_setup_allows_reusing_container_across_apps() -> None:
 
     wireup.integration.starlette.setup(container, app_one)
     wireup.integration.starlette.setup(container, app_two)
+
+
+def write_logs_a(greeter: Injected[GreeterService]) -> None:
+    pass
+
+
+def write_logs_b(greeter: Injected[GreeterService]) -> None:
+    pass
+
+
+def test_wireup_task_caches_regular_functions() -> None:
+    container = wireup.create_async_container(injectables=[shared_services, wireup.integration.starlette])
+    task = WireupTask(container)
+    task._get_injected_wrapper.cache_clear()
+
+    task(write_logs_a)
+    task(write_logs_a)  # should hit cache
+    task(write_logs_b)
+
+    info = task._get_injected_wrapper.cache_info()
+    assert info.hits == 1
+    assert info.misses == 2
+    assert info.currsize == 2
+
+
+def test_wireup_task_does_not_cache_closures() -> None:
+    def make_closure():
+        def write_logs(greeter: Injected[GreeterService]) -> None:
+            pass
+
+        return write_logs
+
+    closure_fn = make_closure()
+    assert "<locals>" in closure_fn.__qualname__
+
+    container = wireup.create_async_container(injectables=[shared_services, wireup.integration.starlette])
+    task = WireupTask(container)
+    task._get_injected_wrapper.cache_clear()
+
+    task(closure_fn)
+    task(closure_fn)
+
+    info = task._get_injected_wrapper.cache_info()
+    assert info.hits == 0
+    assert info.misses == 0
+
+
+def test_wireup_task_does_not_cache_callable_instances() -> None:
+    class CallableTask:
+        def __call__(self, greeter: Injected[GreeterService]) -> None:
+            pass
+
+    callable_instance = CallableTask()
+    assert not isinstance(callable_instance, types.FunctionType)
+
+    container = wireup.create_async_container(injectables=[shared_services, wireup.integration.starlette])
+    task = WireupTask(container)
+    task._get_injected_wrapper.cache_clear()
+
+    task(callable_instance)
+    task(callable_instance)
+
+    info = task._get_injected_wrapper.cache_info()
+    assert info.hits == 0
+    assert info.misses == 0

wireup/integration/starlette.py

@@ -1,5 +1,7 @@
 import contextlib
+import types
 from contextvars import ContextVar
+from functools import lru_cache
 from typing import Any, AsyncIterator
 
 from starlette.applications import Starlette
@@ -113,12 +115,16 @@ def get_request_container() -> ScopedAsyncContainer:
 
 
 class WireupTask:
-    __slots__ = ("container",)
+    __slots__ = ("_get_injected_wrapper", "container")
 
     def __init__(self, container: AsyncContainer) -> None:
         self.container = container
+        self._get_injected_wrapper = lru_cache(maxsize=128)(inject_from_container(self.container))
 
     def __call__(self, fn: AnyCallable) -> Any:
+        should_cache = isinstance(fn, types.FunctionType) and "<locals>" not in fn.__qualname__
+        if should_cache:
+            return self._get_injected_wrapper(fn)
         return inject_from_container(self.container)(fn)