Review feedback: overloads, naming, and test organization

- Add @overload signatures to ConcurrencyLimit.__init__ to clarify the
  three calling conventions (int shorthand, str argument name, keyword-only)
- Rename abbreviated variables (deps, param_name, dep) to full names
- Rename test helper to my_side_effect for clarity
- Move concurrency-specific tests into tests/concurrency_limits/test_annotated.py
  alongside the existing concurrency test suite
- Remove section header comments from tests

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: chrisguidry

src/docket/dependencies/_concurrency.py

@@ -5,7 +5,7 @@
 import asyncio
 import logging
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, overload
 
 from .._cancellation import CANCEL_MSG_CLEANUP, cancel_task
 from ._base import (
@@ -70,21 +70,40 @@ async def expensive(
 
     single: bool = True
 
+    @overload
+    def __init__(
+        self,
+        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:
+        """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:
-        """
-        Args:
-            argument_name: The name of the task argument to use for concurrency grouping.
-                If an ``int`` is passed as the first positional arg, it is treated as
-                *max_concurrent* (convenient for ``Annotated[int, ConcurrencyLimit(1)]``).
-                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)
-        """
         if isinstance(argument_name, int):
             self.argument_name: str | None = None
             self.max_concurrent: int = argument_name
@@ -99,10 +118,9 @@ def __init__(
 
     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=self.argument_name
-            if self.argument_name is not None
-            else name,
+            argument_name,
             max_concurrent=self.max_concurrent,
             scope=self.scope,
         )

src/docket/dependencies/_resolution.py

@@ -29,8 +29,8 @@ def get_single_dependency_parameter_of_type(
     for _, dependency in get_dependency_parameters(function).items():
         if isinstance(dependency, dependency_type):
             return dependency
-    for _, deps in get_annotation_dependencies(function).items():
-        for dependency in deps:
+    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
@@ -86,15 +86,19 @@ async def resolved_dependencies(
                     except Exception as error:
                         arguments[parameter] = FailedDependency(parameter, error)
 
-                annotation_deps = get_annotation_dependencies(execution.function)
-                for param_name, deps in annotation_deps.items():
-                    value = execution.kwargs.get(param_name, arguments.get(param_name))
-                    for dep in deps:
-                        bound = dep.bind_to_parameter(param_name, value)
+                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[param_name] = FailedDependency(param_name, error)
+                            arguments[parameter_name] = FailedDependency(
+                                parameter_name, error
+                            )
 
                 yield arguments
             finally:

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