Add decorator for named parallel branch

Author: Alex Wang

examples/src/parallel/parallel_with_named_branches.py

@@ -1,34 +1,50 @@
-"""Example demonstrating parallel operations with named branches."""
+"""Example demonstrating all parallel branch patterns."""
 
 from typing import Any
 
 from aws_durable_execution_sdk_python.config import ParallelBranch, ParallelConfig
-from aws_durable_execution_sdk_python.context import DurableContext
+from aws_durable_execution_sdk_python.context import (
+    DurableContext,
+    durable_parallel_branch,
+)
 from aws_durable_execution_sdk_python.execution import durable_execution
 
 
+@durable_parallel_branch(name="fetch-orders")
+def fetch_orders(ctx: DurableContext) -> str:
+    return ctx.step(lambda _: "orders-loaded", name="load_orders")
+
+
+@durable_parallel_branch()
+def fetch_preferences(ctx: DurableContext) -> str:
+    return ctx.step(lambda _: "prefs-loaded", name="load_prefs")
+
+
 @durable_execution
 def handler(_event: Any, context: DurableContext) -> list[str]:
-    """Execute named parallel branches using ParallelBranch."""
+    """Execute parallel branches using all supported patterns."""
 
     return context.parallel(
         functions=[
+            # 1. Named parallel branch with ParallelBranch
             ParallelBranch(
                 func=lambda ctx: ctx.step(
                     lambda _: "user-data-loaded", name="load_user"
                 ),
                 name="fetch-user-data",
             ),
+            # 2. Named parallel branch with decorator
+            fetch_orders(),
+            # 3. Unnamed parallel branch with decorator
+            fetch_preferences(),
+            # 4. Unnamed parallel branch with ParallelBranch
             ParallelBranch(
                 func=lambda ctx: ctx.step(
-                    lambda _: "orders-loaded", name="load_orders"
+                    lambda _: "metrics-loaded", name="load_metrics"
                 ),
-                name="fetch-order-history",
-            ),
-            ParallelBranch(
-                func=lambda ctx: ctx.step(lambda _: "prefs-loaded", name="load_prefs"),
-                name="fetch-preferences",
             ),
+            # 5. No wrapper, just a raw callable
+            lambda ctx: ctx.step(lambda _: "config-loaded", name="load_config"),
         ],
         name="load_all_data",
         config=ParallelConfig(max_concurrency=3),

examples/test/parallel/test_parallel_with_named_branches.py

@@ -17,7 +17,7 @@
     lambda_function_name="parallel with named branches",
 )
 def test_parallel_with_named_branches(durable_runner):
-    """Test parallel example with named branches using ParallelBranch."""
+    """Test parallel example with all branch patterns."""
     with durable_runner:
         result = durable_runner.run(input="test", timeout=10)
 
@@ -26,18 +26,30 @@ def test_parallel_with_named_branches(durable_runner):
         "user-data-loaded",
         "orders-loaded",
         "prefs-loaded",
+        "metrics-loaded",
+        "config-loaded",
     ]
 
     # Get the parallel operation
     parallel_op = result.get_context("load_all_data")
     assert parallel_op is not None
     assert parallel_op.status is OperationStatus.SUCCEEDED
 
-    # Verify custom branch names from ParallelBranch
-    assert len(parallel_op.child_operations) == 3
-    child_names = {op.name for op in parallel_op.child_operations}
-    expected_names = {"fetch-user-data", "fetch-order-history", "fetch-preferences"}
-    assert child_names == expected_names
+    # Verify branch names: named branches have custom names, unnamed use defaults
+    assert len(parallel_op.child_operations) == 5
+
+    child_names = [op.name for op in parallel_op.child_operations]
+
+    # 1. Named ParallelBranch
+    assert child_names[0] == "fetch-user-data"
+    # 2. Named decorator
+    assert child_names[1] == "fetch-orders"
+    # 3. Unnamed decorator (None name falls back to index-based default)
+    assert child_names[2] == "parallel-branch-2"
+    # 4. Unnamed ParallelBranch (None name falls back to index-based default)
+    assert child_names[3] == "parallel-branch-3"
+    # 5. Raw callable (no ParallelBranch wrapper, index-based default)
+    assert child_names[4] == "parallel-branch-4"
 
     # Verify all children succeeded
     for child in parallel_op.child_operations:

src/aws_durable_execution_sdk_python/__init__.py

@@ -10,6 +10,7 @@
 from aws_durable_execution_sdk_python.config import ParallelBranch
 from aws_durable_execution_sdk_python.context import (
     DurableContext,
+    durable_parallel_branch,
     durable_step,
     durable_wait_for_callback,
     durable_with_child_context,
@@ -39,6 +40,7 @@
     "ValidationError",
     "__version__",
     "durable_execution",
+    "durable_parallel_branch",
     "durable_step",
     "durable_wait_for_callback",
     "durable_with_child_context",

src/aws_durable_execution_sdk_python/context.py

@@ -121,6 +121,52 @@ def function_with_arguments(child_context: DurableContext):
     return wrapper
 
 
+def durable_parallel_branch(
+    name: str | None = None,
+) -> Callable[
+    [Callable[Concatenate[DurableContext, Params], T]],
+    Callable[Params, ParallelBranch[T]],
+]:
+    """Wrap your callable into a named ParallelBranch for use with context.parallel().
+
+    This is a decorator factory — call it with an optional name to produce
+    the actual decorator.
+
+    Args:
+        name: Optional custom name for this branch. When provided, replaces
+            the default "parallel-branch-{index}" naming in execution history.
+            If None, the function's __name__ is used.
+
+    Example:
+        @durable_parallel_branch(name="fetch-user-data")
+        def fetch_user(ctx: DurableContext, user_id: str) -> dict:
+            return ctx.step(lambda _: {"id": user_id, "name": "Jane"}, name="load_user")
+
+        @durable_parallel_branch(name="fetch-orders")
+        def fetch_orders(ctx: DurableContext, user_id: str) -> list:
+            return ctx.step(lambda _: ["order1", "order2"], name="load_orders")
+
+        # Usage in a durable handler:
+        results = context.parallel(
+            functions=[fetch_user(user_id), fetch_orders(user_id)],
+            name="load-data",
+        )
+    """
+
+    def decorator(
+        func: Callable[Concatenate[DurableContext, Params], T],
+    ) -> Callable[Params, ParallelBranch[T]]:
+        def wrapper(*args, **kwargs) -> ParallelBranch[T]:
+            def function_with_arguments(ctx: DurableContext) -> T:
+                return func(ctx, *args, **kwargs)
+
+            return ParallelBranch(func=function_with_arguments, name=name)
+
+        return wrapper
+
+    return decorator
+
+
 def durable_wait_for_callback(
     func: Callable[Concatenate[str, WaitForCallbackContext, Params], T],
 ) -> Callable[Params, Callable[[str, WaitForCallbackContext], T]]:

tests/context_test.py

@@ -14,13 +14,15 @@
     Duration,
     InvokeConfig,
     MapConfig,
+    ParallelBranch,
     ParallelConfig,
     StepConfig,
 )
 from aws_durable_execution_sdk_python.context import (
     Callback,
     DurableContext,
     ExecutionContext,
+    durable_parallel_branch,
 )
 from aws_durable_execution_sdk_python.exceptions import (
     CallbackError,
@@ -2160,3 +2162,116 @@ def test_should_propagate_outer_parent_id_when_virtual_is_nested_in_virtual():
 
 
 # endregion Virtual-context identity tests
+
+
+# region durable_parallel_branch
+
+
+def test_durable_parallel_branch_returns_parallel_branch_with_name():
+    """Test that the decorator produces a ParallelBranch with the given name."""
+
+    @durable_parallel_branch(name="fetch-user-data")
+    def fetch_user(ctx: DurableContext, user_id: str) -> dict:
+        return {"id": user_id}
+
+    result = fetch_user("user-123")
+
+    assert isinstance(result, ParallelBranch)
+    assert result.name == "fetch-user-data"
+
+
+def test_durable_parallel_branch_with_no_name():
+    """Test that when name is None, ParallelBranch.name is None."""
+
+    @durable_parallel_branch()
+    def fetch_orders(ctx: DurableContext) -> list:
+        return ["order1"]
+
+    result = fetch_orders()
+
+    assert isinstance(result, ParallelBranch)
+    assert result.name is None
+
+
+def test_durable_parallel_branch_callable_delegates_to_func():
+    """Test that calling the ParallelBranch delegates to the wrapped function."""
+
+    @durable_parallel_branch(name="my-branch")
+    def my_branch(ctx: DurableContext, value: int) -> int:
+        return value * 2
+
+    branch = my_branch(21)
+    mock_ctx = Mock(spec=DurableContext)
+
+    result = branch(mock_ctx)
+
+    assert result == 42
+
+
+def test_durable_parallel_branch_with_multiple_args_and_kwargs():
+    """Test that positional and keyword arguments are correctly bound."""
+
+    @durable_parallel_branch(name="compute")
+    def compute(ctx: DurableContext, a: int, b: int, op: str = "add") -> str:
+        if op == "add":
+            return f"{a + b}"
+        return f"{a * b}"
+
+    branch = compute(3, 4, op="mul")
+    mock_ctx = Mock(spec=DurableContext)
+
+    result = branch(mock_ctx)
+
+    assert result == "12"
+
+
+def test_durable_parallel_branch_passes_context_as_first_arg():
+    """Test that the DurableContext is passed as the first argument to the function."""
+    received_ctx = None
+
+    @durable_parallel_branch(name="capture-ctx")
+    def capture(ctx: DurableContext) -> str:
+        nonlocal received_ctx
+        received_ctx = ctx
+        return "done"
+
+    branch = capture()
+    mock_ctx = Mock(spec=DurableContext)
+    branch(mock_ctx)
+
+    assert received_ctx is mock_ctx
+
+
+def test_durable_parallel_branch_multiple_invocations_are_independent():
+    """Test that calling the wrapper multiple times produces independent branches."""
+
+    @durable_parallel_branch(name="greet")
+    def greet(ctx: DurableContext, name: str) -> str:
+        return f"hello {name}"
+
+    branch_a = greet("Alice")
+    branch_b = greet("Bob")
+
+    mock_ctx = Mock(spec=DurableContext)
+
+    assert branch_a(mock_ctx) == "hello Alice"
+    assert branch_b(mock_ctx) == "hello Bob"
+
+
+def test_durable_parallel_branch_is_compatible_with_parallel_functions_arg():
+    """Test that the result can be used in a functions list alongside plain callables."""
+
+    @durable_parallel_branch(name="named-branch")
+    def named(ctx: DurableContext) -> str:
+        return "named"
+
+    plain = lambda ctx: "plain"  # noqa: E731
+
+    functions = [named(), plain]
+
+    assert isinstance(functions[0], ParallelBranch)
+    assert callable(functions[0])
+    assert callable(functions[1])
+
+
+# endregion durable_parallel_branch