44
55from taskbadger .exceptions import (
66 ConfigurationError ,
7+ MissingConfiguration ,
78 ServerError ,
89 TaskbadgerException ,
910 Unauthorized ,
2122 PatchedTaskRequestTags ,
2223 StatusEnum ,
2324 TaskRequest ,
24- TaskRequestTags ,
2525)
2626from taskbadger .internal .types import UNSET
27- from taskbadger .mug import Badger , Session , Settings
27+ from taskbadger .mug import Badger , Callback , Session , Settings
2828from taskbadger .systems import System
29+ from taskbadger .utils import import_string
2930
3031log = logging .getLogger ("taskbadger" )
3132
@@ -38,12 +39,13 @@ def init(
3839 token : str = None ,
3940 systems : list [System ] = None ,
4041 tags : dict [str , str ] = None ,
42+ before_create : Callback = None ,
4143):
4244 """Initialize Task Badger client
4345
4446 Call this function once per thread
4547 """
46- _init (_TB_HOST , organization_slug , project_slug , token , systems , tags )
48+ _init (_TB_HOST , organization_slug , project_slug , token , systems , tags , before_create )
4749
4850
4951def _init (
@@ -53,12 +55,19 @@ def _init(
5355 token : str = None ,
5456 systems : list [System ] = None ,
5557 tags : dict [str , str ] = None ,
58+ before_create : Callback = None ,
5659):
5760 host = host or os .environ .get ("TASKBADGER_HOST" , "https://taskbadger.net" )
5861 organization_slug = organization_slug or os .environ .get ("TASKBADGER_ORG" )
5962 project_slug = project_slug or os .environ .get ("TASKBADGER_PROJECT" )
6063 token = token or os .environ .get ("TASKBADGER_API_KEY" )
6164
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+
6271 if host and organization_slug and project_slug and token :
6372 systems = systems or []
6473 settings = Settings (
@@ -67,10 +76,11 @@ def _init(
6776 organization_slug ,
6877 project_slug ,
6978 systems = {system .identifier : system for system in systems },
79+ before_create = before_create ,
7080 )
7181 Badger .current .bind (settings , tags )
7282 else :
73- raise ConfigurationError (
83+ raise MissingConfiguration (
7484 host = host ,
7585 organization_slug = organization_slug ,
7686 project_slug = project_slug ,
@@ -118,29 +128,33 @@ def create_task(
118128 Returns:
119129 Task: The created Task object.
120130 """
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
135143 scope = Badger .current .scope ()
136144 if scope .context or data :
137145 data = data or {}
138- task . data = {** scope .context , ** data }
146+ task_dict [ " data" ] = {** scope .context , ** data }
139147 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 ]
141149 if scope .tags or tags :
142150 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 )
144158 kwargs = _make_args (body = task )
145159 if monitor_id :
146160 kwargs ["x_taskbadger_monitor" ] = monitor_id
0 commit comments