Support Annotated-style dependency injection (#354)

chrisguidry · claude · web-flow · commit 7180b799d187 · 2026-02-27T14:50:10.000-05:00
## Summary Dependencies can now be attached as `Annotated` type-hint metadata instead of only as default parameter values: ```python async def process(customer_id: Annotated[int, ConcurrencyLimit(1)]): ... ``` The parameter keeps its real value; the dependency runs as a side-effect. This is especially nice for ConcurrencyLimit where the old default-param style required a separate dummy parameter just to carry the dependency. Changes: - `resolved_dependencies()` now calls `get_annotation_dependencies()` from uncalled-for and enters each annotation dep via `bind_to_parameter()` - `ConcurrencyLimit` gains `bind_to_parameter()` to auto-infer the argument name from the annotated parameter - `ConcurrencyLimit(1)` shorthand: passing an int as the first positional arg sets `max_concurrent` (convenient for the Annotated style) - `get_single_dependency_parameter_of_type()` now searches annotations too - Bumps uncalled-for to >=0.2.0 for the annotation extraction API Includes contract tests for the uncalled-for behaviors we depend on, and integration tests covering concurrency limits, Depends side-effects, mixed styles, type aliases, and validation. Closes #334, closes #163. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
diff --git a/pyproject.toml b/pyproject.toml
@@ -37,7 +37,7 @@ dependencies = [
     "typer>=0.15.1",
     "typing_extensions>=4.12.0",
     "tzdata>=2025.2; sys_platform == 'win32'",
-    "uncalled-for>=0.1.2",
+    "uncalled-for>=0.2.0",
 ]
 
 [project.optional-dependencies]
diff --git a/src/docket/dependencies/__init__.py b/src/docket/dependencies/__init__.py
@@ -41,6 +41,7 @@
 from ._progress import Progress
 from ._resolution import (
     FailedDependency,
+    get_annotation_dependencies,
     get_single_dependency_of_type,
     get_single_dependency_parameter_of_type,
     resolved_dependencies,
@@ -72,6 +73,7 @@
     "DependencyFunction",
     "Shared",
     "SharedContext",
+    "get_annotation_dependencies",
     "get_dependency_parameters",
     # Retry
     "ForcedRetry",
diff --git a/src/docket/dependencies/_concurrency.py b/src/docket/dependencies/_concurrency.py
@@ -5,7 +5,7 @@
 import asyncio
 import logging
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, overload
 
 from .._cancellation import CANCEL_MSG_CLEANUP, cancel_task
 from ._base import (
@@ -49,54 +49,82 @@ class ConcurrencyLimit(Dependency["ConcurrencyLimit"]):
 
     Can limit concurrency globally for a task, or per specific argument value.
 
-    Example:
+    Works both as a default parameter and as ``Annotated`` metadata::
 
-    ```python
-    async def expensive_operation(
-        concurrency: ConcurrencyLimit = ConcurrencyLimit(max_concurrent=3)
-    ) -> None:
-        # Only 3 instances of this task will run at a time
-        ...
+        # Default-parameter style
+        async def process_customer(
+            customer_id: int,
+            concurrency: ConcurrencyLimit = ConcurrencyLimit("customer_id", 1),
+        ) -> None: ...
 
-    async def process_customer(
-        customer_id: int,
-        concurrency: ConcurrencyLimit = ConcurrencyLimit("customer_id", max_concurrent=1)
-    ) -> None:
-        # Only one task per customer_id will run at a time
-        ...
+        # Annotated style (parameter name auto-inferred)
+        async def process_customer(
+            customer_id: Annotated[int, ConcurrencyLimit(1)],
+        ) -> None: ...
 
-    async def backup_db(
-        db_name: str,
-        concurrency: ConcurrencyLimit = ConcurrencyLimit("db_name", max_concurrent=3)
-    ) -> None:
-        # Only 3 backup tasks per database name will run at a time
-        ...
-    ```
+        # Per-task (no argument grouping)
+        async def expensive(
+            concurrency: ConcurrencyLimit = ConcurrencyLimit(max_concurrent=3),
+        ) -> None: ...
     """
 
     single: bool = True
 
+    @overload
     def __init__(
         self,
-        argument_name: str | None = None,
+        max_concurrent: int,
+        /,
+        *,
+        scope: str | None = None,
+    ) -> None:
+        """Annotated style: ``Annotated[int, ConcurrencyLimit(1)]``."""
+
+    @overload
+    def __init__(
+        self,
+        argument_name: str,
         max_concurrent: int = 1,
         scope: str | None = None,
     ) -> None:
-        """
-        Args:
-            argument_name: The name of the task argument to use for concurrency grouping.
-                If None, limits concurrency for the task function itself.
-            max_concurrent: Maximum number of concurrent tasks
-            scope: Optional scope prefix for Redis keys (defaults to docket name)
-        """
-        self.argument_name = argument_name
-        self.max_concurrent = max_concurrent
+        """Default-param style with per-argument grouping."""
+
+    @overload
+    def __init__(
+        self,
+        *,
+        max_concurrent: int = 1,
+        scope: str | None = None,
+    ) -> None:
+        """Per-task concurrency (no argument grouping)."""
+
+    def __init__(
+        self,
+        argument_name: str | int | None = None,
+        max_concurrent: int = 1,
+        scope: str | None = None,
+    ) -> None:
+        if isinstance(argument_name, int):
+            self.argument_name: str | None = None
+            self.max_concurrent: int = argument_name
+        else:
+            self.argument_name = argument_name
+            self.max_concurrent = max_concurrent
         self.scope = scope
         self._concurrency_key: str | None = None
         self._initialized: bool = False
         self._task_key: str | None = None
         self._renewal_task: asyncio.Task[None] | None = None
 
+    def bind_to_parameter(self, name: str, value: Any) -> ConcurrencyLimit:
+        """Bind to an ``Annotated`` parameter, inferring argument_name if not set."""
+        argument_name = self.argument_name if self.argument_name is not None else name
+        return ConcurrencyLimit(
+            argument_name,
+            max_concurrent=self.max_concurrent,
+            scope=self.scope,
+        )
+
     async def __aenter__(self) -> ConcurrencyLimit:
         from ._functional import _Depends
 
diff --git a/src/docket/dependencies/_functional.py b/src/docket/dependencies/_functional.py
@@ -9,7 +9,9 @@
     DependencyFactory,
     Shared as Shared,
     SharedContext as SharedContext,
-    _Depends as _UncalledForDepends,
+)
+from uncalled_for.functional import _Depends as _UncalledForDepends
+from uncalled_for.introspection import (
     _parameter_cache as _parameter_cache,
     get_dependency_parameters,
 )
diff --git a/src/docket/dependencies/_resolution.py b/src/docket/dependencies/_resolution.py
@@ -7,6 +7,7 @@
 
 from uncalled_for import (
     FailedDependency as FailedDependency,
+    get_annotation_dependencies as get_annotation_dependencies,
     validate_dependencies as validate_dependencies,
 )
 
@@ -28,6 +29,10 @@ def get_single_dependency_parameter_of_type(
     for _, dependency in get_dependency_parameters(function).items():
         if isinstance(dependency, dependency_type):
             return dependency
+    for _, dependencies in get_annotation_dependencies(function).items():
+        for dependency in dependencies:
+            if isinstance(dependency, dependency_type):
+                return dependency  # type: ignore[return-value]
     return None
 
 
@@ -81,6 +86,20 @@ async def resolved_dependencies(
                     except Exception as error:
                         arguments[parameter] = FailedDependency(parameter, error)
 
+                annotations = get_annotation_dependencies(execution.function)
+                for parameter_name, dependencies in annotations.items():
+                    value = execution.kwargs.get(
+                        parameter_name, arguments.get(parameter_name)
+                    )
+                    for dependency in dependencies:
+                        bound = dependency.bind_to_parameter(parameter_name, value)
+                        try:
+                            await stack.enter_async_context(bound)
+                        except Exception as error:
+                            arguments[parameter_name] = FailedDependency(
+                                parameter_name, error
+                            )
+
                 yield arguments
             finally:
                 _Depends.stack.reset(stack_token)
diff --git a/tests/concurrency_limits/test_annotated.py b/tests/concurrency_limits/test_annotated.py
@@ -0,0 +1,117 @@
+"""Tests for ConcurrencyLimit via Annotated-style type hints."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import Annotated
+
+import pytest
+
+from docket import ConcurrencyLimit, Docket, Worker
+
+from tests.concurrency_limits.overlap import assert_some_overlap
+
+
+async def test_annotated_concurrency_limit(docket: Docket, worker: Worker):
+    """Annotated[int, ConcurrencyLimit(1)] limits concurrency per-parameter."""
+    results: list[str] = []
+
+    async def task(customer_id: Annotated[int, ConcurrencyLimit(1)]):
+        results.append(f"start_{customer_id}")
+        await asyncio.sleep(0.01)
+        results.append(f"end_{customer_id}")
+
+    await docket.add(task)(customer_id=1)
+    await docket.add(task)(customer_id=1)
+
+    await worker.run_until_finished()
+
+    assert results == ["start_1", "end_1", "start_1", "end_1"]
+
+
+async def test_annotated_concurrency_different_values(docket: Docket, worker: Worker):
+    """Different argument values get independent concurrency slots."""
+    execution_intervals: dict[int, tuple[float, float]] = {}
+
+    async def task(customer_id: Annotated[int, ConcurrencyLimit(1)]):
+        start = time.monotonic()
+        await asyncio.sleep(0.1)
+        end = time.monotonic()
+        execution_intervals[customer_id] = (start, end)
+
+    await docket.add(task)(customer_id=1)
+    await docket.add(task)(customer_id=2)
+    await docket.add(task)(customer_id=3)
+
+    worker.concurrency = 10
+    await worker.run_until_finished()
+
+    assert len(execution_intervals) == 3
+    intervals = list(execution_intervals.values())
+    assert_some_overlap(intervals, "Different customers should run concurrently")
+
+
+async def test_annotated_concurrency_max_concurrent(docket: Docket, worker: Worker):
+    """max_concurrent>1 allows that many concurrent executions per value."""
+    active_tasks: list[int] = []
+    max_concurrent_seen = 0
+    lock = asyncio.Lock()
+
+    async def task(
+        db_name: str,
+        task_id: Annotated[int, ConcurrencyLimit("db_name", max_concurrent=2)],
+    ):
+        nonlocal max_concurrent_seen
+        async with lock:
+            active_tasks.append(task_id)
+            max_concurrent_seen = max(max_concurrent_seen, len(active_tasks))
+        await asyncio.sleep(0.1)
+        async with lock:
+            active_tasks.remove(task_id)
+
+    for i in range(5):
+        await docket.add(task)(db_name="postgres", task_id=i)
+
+    worker.concurrency = 10
+    await worker.run_until_finished()
+
+    assert max_concurrent_seen <= 2
+
+
+async def test_two_annotated_concurrency_limits_rejected(docket: Docket):
+    """ConcurrencyLimit.single=True prevents two annotations on one function."""
+    with pytest.raises(ValueError, match="Only one ConcurrencyLimit"):
+
+        async def task(
+            customer_id: Annotated[int, ConcurrencyLimit(1)],
+            region: Annotated[str, ConcurrencyLimit(2)],
+        ): ...  # pragma: no cover
+
+        await docket.add(task)(customer_id=1, region="us")
+
+
+async def test_single_conflict_annotation_and_default(docket: Docket):
+    """single=True conflict detected across annotation and default-param styles."""
+    with pytest.raises(ValueError, match="Only one ConcurrencyLimit"):
+
+        async def task(
+            customer_id: Annotated[int, ConcurrencyLimit(1)],
+            concurrency: ConcurrencyLimit = ConcurrencyLimit(max_concurrent=2),
+        ): ...  # pragma: no cover
+
+        await docket.add(task)(customer_id=1)
+
+
+async def test_annotated_concurrency_keys_cleaned_up(docket: Docket, worker: Worker):
+    """Concurrency keys from annotated deps are properly cleaned up."""
+
+    async def task(customer_id: Annotated[int, ConcurrencyLimit(1)]):
+        pass
+
+    await docket.add(task)(customer_id=42)
+    await worker.run_until_finished()
+
+    async with docket.redis() as redis:
+        key = f"{docket.name}:concurrency:customer_id:42"
+        assert await redis.exists(key) == 0
diff --git a/tests/test_annotated_dependencies.py b/tests/test_annotated_dependencies.py
diff --git a/tests/test_uncalled_for_contract.py b/tests/test_uncalled_for_contract.py
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -37,7 +37,7 @@ dependencies = [`
`37`	`37`	`"typer>=0.15.1",`
`38`	`38`	`"typing_extensions>=4.12.0",`
`39`	`39`	`"tzdata>=2025.2; sys_platform == 'win32'",`
`40`		`- "uncalled-for>=0.1.2",`
	`40`	`+ "uncalled-for>=0.2.0",`
`41`	`41`	`]`
`42`	`42`
`43`	`43`	`[project.optional-dependencies]`