Support *args and **kwargs in add_execution

Author: edtechre

src/pybroker/strategy.py

@@ -102,13 +102,17 @@ class Execution(NamedTuple):
             execution of ``fn``.
         indicator_names: Names of :class:`pybroker.indicator.Indicator`\ s
             used for execution of ``fn``.
+        args: Additional positional arguments for ``fn``.
+        kwargs: Additional keyword arguments for ``fn``.
     """
 
     id: int
     symbols: frozenset[str]
     fn: Optional[Callable[[ExecContext], None]]
     model_names: frozenset[str]
     indicator_names: frozenset[str]
+    args: tuple[Any, ...] = tuple()
+    kwargs: tuple[tuple[str, Any], ...] = tuple()
 
 
 class BacktestMixin:
@@ -186,6 +190,8 @@ def backtest_executions(
         pending_order_scope = PendingOrderScope()
         exec_ctxs: dict[str, ExecContext] = {}
         exec_fns: dict[str, Callable[[ExecContext], None]] = {}
+        exec_args: dict[str, tuple[Any, ...]] = {}
+        exec_kwargs: dict[str, tuple[tuple[str, Any], ...]] = {}
         for sym in test_syms:
             for exec in executions:
                 if sym not in exec.symbols:
@@ -203,6 +209,8 @@ def backtest_executions(
                     sym_end_index=sym_end_index,
                     session=sessions[sym],
                 )
+                exec_args[sym] = exec.args
+                exec_kwargs[sym] = exec.kwargs
                 if exec.fn is not None:
                     exec_fns[sym] = exec.fn
         sym_exec_dates = {
@@ -311,7 +319,11 @@ def backtest_executions(
                 before_exec_fn(active_ctxs)
             for sym, ctx in active_ctxs.items():
                 if sym in exec_fns:
-                    exec_fns[sym](ctx)
+                    exec_fns[sym](
+                        ctx,
+                        *exec_args.get(sym, ()),
+                        **dict(exec_kwargs.get(sym, ())),
+                    )
             if after_exec_fn is not None and active_ctxs:
                 after_exec_fn(active_ctxs)
             for ctx in active_ctxs.values():
@@ -896,6 +908,8 @@ def add_execution(
         symbols: Union[str, Iterable[str]],
         models: Optional[Union[ModelSource, Iterable[ModelSource]]] = None,
         indicators: Optional[Union[Indicator, Iterable[Indicator]]] = None,
+        *args: Any,
+        **kwargs: Any,
     ):
         r"""Adds an execution to backtest.
 
@@ -910,6 +924,8 @@ def add_execution(
             indicators: :class:`Iterable` of
                 :class:`pybroker.indicator.Indicator`\ s to compute for
                 backtesting.
+            args: Positional arguments passed to ``fn``.
+            kwargs: Keyword arguments passed to ``fn``.
         """
         symbols = (
             frozenset((symbols,))
@@ -978,6 +994,8 @@ def add_execution(
                 fn=fn,
                 model_names=model_names,
                 indicator_names=ind_names,
+                args=args,
+                kwargs=tuple(sorted(kwargs.items())),
             )
         )
 

tests/test_strategy.py

@@ -2515,3 +2515,31 @@ def exec_fn(ctx):
         strategy.add_execution(exec_fn, "SPY")
         with pytest.raises(ValueError, match=re.escape("warmup must be > 0.")):
             strategy.backtest(warmup=-1)
+
+    def test_backtest_when_args_and_kwargs(self, data_source_df):
+        def exec_fn(ctx, foo, bar=None):
+            assert foo == "foo_value"
+            assert bar == "bar_value"
+
+        strategy = Strategy(data_source_df, START_DATE, END_DATE)
+        strategy.add_execution(
+            exec_fn,
+            "SPY",
+            foo="foo_value",
+            bar="bar_value",
+        )
+        strategy.backtest(calc_bootstrap=False)
+
+    def test_walkforward_when_args_and_kwargs(self, data_source_df):
+        def exec_fn(ctx, foo, bar=None):
+            assert foo == "foo_value"
+            assert bar == "bar_value"
+
+        strategy = Strategy(data_source_df, START_DATE, END_DATE)
+        strategy.add_execution(
+            exec_fn,
+            "SPY",
+            foo="foo_value",
+            bar="bar_value",
+        )
+        strategy.walkforward(windows=2, calc_bootstrap=False)