Coverage for src/shephex/decorators/chain.py: 100%
21 statements
« prev ^ index » next coverage.py v7.6.1, created at 2025-06-20 14:13 +0200
« prev ^ index » next coverage.py v7.6.1, created at 2025-06-20 14:13 +0200
1import functools
2from pathlib import Path
3from typing import Callable, Literal, Optional
5from shephex.decorators import get_decorator_state
6from shephex.experiment import ChainableExperimentIterator
7from shephex.experiment.procedure import PickleProcedure, ScriptProcedure
10def chain(
11 context: bool = False,
12 hex_directory: str | Path = 'experiments/',
13 procedure_type: Literal['ScriptProcedure', 'PickleProcedure'] = 'PickleProcedure',
14) -> Callable:
15 """
16 Decorator to create a ChainableExperimentIterator for the function it is applied to.
18 Parameters
19 -----------
20 context (bool):
21 Whether to pass the experiment context to the function. If this is activated
22 the function will need to take an additional keyword argument `context` that
23 will be passed the experiment context. This offers limited access to
24 writing to the context, e.g. to track progress.
25 """
27 def decorator(function: Callable):
29 @functools.wraps(function)
30 def wrapper(*args, directory: Optional[str | Path] = None, **kwargs):
32 # If the decorator is not active, return the function
33 if not get_decorator_state().active:
34 return function
36 if procedure_type == 'ScriptProcedure':
37 procedure = ScriptProcedure(function, context=context)
38 elif procedure_type == 'PickleProcedure':
39 procedure = PickleProcedure(function, context=context)
41 if directory is None:
42 directory = hex_directory
44 return ChainableExperimentIterator(procedure, directory)
46 return wrapper
48 return decorator