Coverage for src/shephex/decorators/hexperiment.py: 100%

22 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2025-03-29 18:45 +0100

1import functools 

2from pathlib import Path 

3from typing import Callable, Literal, Optional 

4 

5from shephex.decorators import get_decorator_state 

6from shephex.experiment import Experiment 

7from shephex.experiment.procedure import PickleProcedure, ScriptProcedure 

8 

9 

10def hexperiment( 

11 context: bool = False, 

12 hex_directory: str | Path = 'experiments/', 

13 procedure_type: Literal['ScriptProcedure', 'PickleProcedure'] = 'PickleProcedure', 

14) -> Callable: 

15 """ 

16 Decorator to create an Experiment 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 hex_directory (str | Path): 

26 The directory to save the experiment to. This can be a string or a Path object. 

27 procedure_type (Literal['ScriptProcedure', 'PickleProcedure']): 

28 Type of procedure used, by default uses ScriptProcedure such that no  

29 pickling is required. 

30 

31 Returns 

32 -------- 

33 decorator: Callable 

34 A decorator function that takes the function to be decorated and returns an 

35 Experiment that can 

36 """ 

37 

38 def decorator(function: Callable): 

39 

40 @functools.wraps(function) 

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

42 

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

44 if not get_decorator_state().active: 

45 return function 

46 

47 # If decorator is active, make an experiment 

48 if procedure_type == 'ScriptProcedure': 

49 procedure = ScriptProcedure(function, context=context) 

50 elif procedure_type == 'PickleProcedure': 

51 procedure = PickleProcedure(function, context=context) 

52 

53 if directory is None: 

54 directory = hex_directory 

55 

56 experiment = Experiment( 

57 *args, procedure=procedure, root_path=directory, **kwargs) 

58 

59 return experiment 

60 

61 return wrapper 

62 

63 return decorator