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

1import functools 

2from pathlib import Path 

3from typing import Callable, Literal, Optional 

4 

5from shephex.decorators import get_decorator_state 

6from shephex.experiment import ChainableExperimentIterator 

7from shephex.experiment.procedure import PickleProcedure, ScriptProcedure 

8 

9 

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. 

17 

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 """ 

26 

27 def decorator(function: Callable): 

28 

29 @functools.wraps(function) 

30 def wrapper(*args, directory: Optional[str | Path] = None, **kwargs): 

31 

32 # If the decorator is not active, return the function 

33 if not get_decorator_state().active: 

34 return function 

35 

36 if procedure_type == 'ScriptProcedure': 

37 procedure = ScriptProcedure(function, context=context) 

38 elif procedure_type == 'PickleProcedure': 

39 procedure = PickleProcedure(function, context=context) 

40 

41 if directory is None: 

42 directory = hex_directory 

43 

44 return ChainableExperimentIterator(procedure, directory) 

45 

46 return wrapper 

47 

48 return decorator