Coverage for /Users/jerry/Development/yenta/yenta/tasks/Task.py: 100%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

59 statements  

1from dataclasses import dataclass, field 

2from enum import Enum 

3from functools import wraps 

4from inspect import signature 

5from typing import Callable, List, Dict, Optional 

6 

7 

8class ParameterType(int, Enum): 

9 

10 PIPELINE_RESULTS = 1 

11 EXPLICIT = 2 

12 

13 

14class ResultType(str, Enum): 

15 

16 VALUE = 'values' 

17 ARTIFACT = 'artifacts' 

18 

19 

20@dataclass 

21class ResultSpec: 

22 

23 result_task_name: str 

24 result_var_name: Optional[str] = None 

25 

26 

27@dataclass 

28class ParameterSpec: 

29 

30 param_name: str 

31 param_type: ParameterType 

32 result_spec: Optional[ResultSpec] = None 

33 selector: Optional[Callable] = None 

34 

35 

36@dataclass 

37class TaskDef: 

38 

39 name: str 

40 depends_on: Optional[List[str]] 

41 pure: bool 

42 param_specs: List[ParameterSpec] = field(default_factory=list) 

43 

44 

45class InvalidTaskDefinitionError(Exception): 

46 pass 

47 

48 

49def build_parameter_spec(func, depends_on: List[str]): 

50 

51 sig = signature(func) 

52 param_names = list(sig.parameters.keys()) 

53 

54 if len(param_names) < len(depends_on): 

55 

56 raise InvalidTaskDefinitionError( 

57 f'Insufficient number of parameters ({len(param_names)}) defined ' 

58 f'for a task that depends on {len(depends_on)} inputs.') 

59 elif len(param_names) == 0: 

60 spec = [] 

61 else: 

62 spec = [] 

63 for dependency, param_name in zip(depends_on, param_names): 

64 if '.' not in dependency: 

65 task_dep_name = dependency 

66 var_dep_name = None 

67 param_type = ParameterType.PIPELINE_RESULTS 

68 else: 

69 task_dep_name, var_dep_name = dependency.split('.') 

70 param_type = ParameterType.EXPLICIT 

71 result_spec = ResultSpec(task_dep_name, var_dep_name) 

72 spec.append(ParameterSpec(param_name, param_type, result_spec)) 

73 

74 return spec 

75 

76 

77def task(_func=None, *, depends_on: Optional[List[str]] = None, pure: bool = True, selectors=None): 

78 

79 depends_on = depends_on or [] 

80 

81 def decorator_task(func: Callable): 

82 

83 @wraps(func) 

84 def task_wrapper(*args, **kwargs): 

85 return func(*args, **kwargs) 

86 

87 setattr(task_wrapper, 'task_def', TaskDef( 

88 name=func.__name__, 

89 depends_on=depends_on, 

90 pure=pure, 

91 param_specs=build_parameter_spec(func, depends_on) 

92 )) 

93 

94 setattr(task_wrapper, '_yenta_task', True) 

95 

96 return task_wrapper 

97 

98 if _func is None: 

99 return decorator_task 

100 else: 

101 return decorator_task(_func)