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
« 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
5from shephex.decorators import get_decorator_state
6from shephex.experiment import Experiment
7from shephex.experiment.procedure import PickleProcedure, ScriptProcedure
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.
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.
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 """
38 def decorator(function: Callable):
40 @functools.wraps(function)
41 def wrapper(*args, directory: Optional[str | Path] = None, **kwargs):
43 # If the decorator is not active, return the function
44 if not get_decorator_state().active:
45 return function
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)
53 if directory is None:
54 directory = hex_directory
56 experiment = Experiment(
57 *args, procedure=procedure, root_path=directory, **kwargs)
59 return experiment
61 return wrapper
63 return decorator