4
4
5
5
from taskbadger .exceptions import (
6
6
ConfigurationError ,
7
+ MissingConfiguration ,
7
8
ServerError ,
8
9
TaskbadgerException ,
9
10
Unauthorized ,
21
22
PatchedTaskRequestTags ,
22
23
StatusEnum ,
23
24
TaskRequest ,
24
- TaskRequestTags ,
25
25
)
26
26
from taskbadger .internal .types import UNSET
27
- from taskbadger .mug import Badger , Session , Settings
27
+ from taskbadger .mug import Badger , Callback , Session , Settings
28
28
from taskbadger .systems import System
29
+ from taskbadger .utils import import_string
29
30
30
31
log = logging .getLogger ("taskbadger" )
31
32
@@ -38,12 +39,13 @@ def init(
38
39
token : str = None ,
39
40
systems : list [System ] = None ,
40
41
tags : dict [str , str ] = None ,
42
+ before_create : Callback = None ,
41
43
):
42
44
"""Initialize Task Badger client
43
45
44
46
Call this function once per thread
45
47
"""
46
- _init (_TB_HOST , organization_slug , project_slug , token , systems , tags )
48
+ _init (_TB_HOST , organization_slug , project_slug , token , systems , tags , before_create )
47
49
48
50
49
51
def _init (
@@ -53,12 +55,19 @@ def _init(
53
55
token : str = None ,
54
56
systems : list [System ] = None ,
55
57
tags : dict [str , str ] = None ,
58
+ before_create : Callback = None ,
56
59
):
57
60
host = host or os .environ .get ("TASKBADGER_HOST" , "https://taskbadger.net" )
58
61
organization_slug = organization_slug or os .environ .get ("TASKBADGER_ORG" )
59
62
project_slug = project_slug or os .environ .get ("TASKBADGER_PROJECT" )
60
63
token = token or os .environ .get ("TASKBADGER_API_KEY" )
61
64
65
+ if before_create and isinstance (before_create , str ):
66
+ try :
67
+ before_create = import_string (before_create )
68
+ except ImportError as e :
69
+ raise ConfigurationError (f"Could not import module: { before_create } " ) from e
70
+
62
71
if host and organization_slug and project_slug and token :
63
72
systems = systems or []
64
73
settings = Settings (
@@ -67,10 +76,11 @@ def _init(
67
76
organization_slug ,
68
77
project_slug ,
69
78
systems = {system .identifier : system for system in systems },
79
+ before_create = before_create ,
70
80
)
71
81
Badger .current .bind (settings , tags )
72
82
else :
73
- raise ConfigurationError (
83
+ raise MissingConfiguration (
74
84
host = host ,
75
85
organization_slug = organization_slug ,
76
86
project_slug = project_slug ,
@@ -118,29 +128,33 @@ def create_task(
118
128
Returns:
119
129
Task: The created Task object.
120
130
"""
121
- value = _none_to_unset (value )
122
- value_max = _none_to_unset (value_max )
123
- data = _none_to_unset (data )
124
- max_runtime = _none_to_unset (max_runtime )
125
- stale_timeout = _none_to_unset (stale_timeout )
126
-
127
- task = TaskRequest (
128
- name = name ,
129
- status = status ,
130
- value = value ,
131
- value_max = value_max ,
132
- max_runtime = max_runtime ,
133
- stale_timeout = stale_timeout ,
134
- )
131
+ task_dict = {
132
+ "name" : name ,
133
+ "status" : status ,
134
+ }
135
+ if value is not None :
136
+ task_dict ["value" ] = value
137
+ if value_max is not None :
138
+ task_dict ["value_max" ] = value_max
139
+ if max_runtime is not None :
140
+ task_dict ["max_runtime" ] = max_runtime
141
+ if stale_timeout is not None :
142
+ task_dict ["stale_timeout" ] = stale_timeout
135
143
scope = Badger .current .scope ()
136
144
if scope .context or data :
137
145
data = data or {}
138
- task . data = {** scope .context , ** data }
146
+ task_dict [ " data" ] = {** scope .context , ** data }
139
147
if actions :
140
- task . additional_properties = { "actions" : [a .to_dict () for a in actions ]}
148
+ task_dict [ "actions" ] = [a .to_dict () for a in actions ]
141
149
if scope .tags or tags :
142
150
tags = tags or {}
143
- task .tags = TaskRequestTags .from_dict ({** scope .tags , ** tags })
151
+ task_dict ["tags" ] = {** scope .tags , ** tags }
152
+
153
+ task_dict = Badger .current .call_before_create (task_dict )
154
+ if not task_dict :
155
+ raise TaskbadgerException ("before_create callback returned None" )
156
+
157
+ task = TaskRequest .from_dict (task_dict )
144
158
kwargs = _make_args (body = task )
145
159
if monitor_id :
146
160
kwargs ["x_taskbadger_monitor" ] = monitor_id
0 commit comments