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
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
1from dataclasses import dataclass, field
2from enum import Enum
3from functools import wraps
4from inspect import signature
5from typing import Callable, List, Dict, Optional
8class ParameterType(int, Enum):
10 PIPELINE_RESULTS = 1
11 EXPLICIT = 2
14class ResultType(str, Enum):
16 VALUE = 'values'
17 ARTIFACT = 'artifacts'
20@dataclass
21class ResultSpec:
23 result_task_name: str
24 result_var_name: Optional[str] = None
27@dataclass
28class ParameterSpec:
30 param_name: str
31 param_type: ParameterType
32 result_spec: Optional[ResultSpec] = None
33 selector: Optional[Callable] = None
36@dataclass
37class TaskDef:
39 name: str
40 depends_on: Optional[List[str]]
41 pure: bool
42 param_specs: List[ParameterSpec] = field(default_factory=list)
45class InvalidTaskDefinitionError(Exception):
46 pass
49def build_parameter_spec(func, depends_on: List[str]):
51 sig = signature(func)
52 param_names = list(sig.parameters.keys())
54 if len(param_names) < len(depends_on):
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))
74 return spec
77def task(_func=None, *, depends_on: Optional[List[str]] = None, pure: bool = True, selectors=None):
79 depends_on = depends_on or []
81 def decorator_task(func: Callable):
83 @wraps(func)
84 def task_wrapper(*args, **kwargs):
85 return func(*args, **kwargs)
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 ))
94 setattr(task_wrapper, '_yenta_task', True)
96 return task_wrapper
98 if _func is None:
99 return decorator_task
100 else:
101 return decorator_task(_func)