Skip to content

Commit ed1bc4e

Browse files
authored
Merge pull request #34 from taskbadger/sk/fix-scope
add 'before_create' callback
2 parents dbffcf1 + 195e0fb commit ed1bc4e

File tree

6 files changed

+131
-23
lines changed

6 files changed

+131
-23
lines changed

Diff for: taskbadger/exceptions.py

+4
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,8 @@
11
class ConfigurationError(Exception):
2+
pass
3+
4+
5+
class MissingConfiguration(ConfigurationError):
26
def __init__(self, **kwargs):
37
self.missing = [name for name, arg in kwargs.items() if arg is None]
48

Diff for: taskbadger/mug.py

+10-1
Original file line numberDiff line numberDiff line change
@@ -2,21 +2,25 @@
22
from contextlib import ContextDecorator
33
from contextvars import ContextVar
44
from copy import deepcopy
5-
from typing import Union
5+
from typing import Callable, Optional, Union
66

77
from taskbadger.internal import AuthenticatedClient
88
from taskbadger.systems import System
99

1010
_local = ContextVar("taskbadger_client")
1111

1212

13+
Callback = Union[str, Callable[[dict], Optional[dict]]]
14+
15+
1316
@dataclasses.dataclass
1417
class Settings:
1518
base_url: str
1619
token: str
1720
organization_slug: str
1821
project_slug: str
1922
systems: dict[str, System] = dataclasses.field(default_factory=dict)
23+
before_create: Callback = None
2024

2125
def get_client(self):
2226
return AuthenticatedClient(self.base_url, self.token)
@@ -140,6 +144,11 @@ def client(self) -> AuthenticatedClient:
140144
def scope(self) -> Scope:
141145
return self._scope
142146

147+
def call_before_create(self, task: dict) -> Optional[dict]:
148+
if self.settings and self.settings.before_create:
149+
return self.settings.before_create(task)
150+
return task
151+
143152
@classmethod
144153
def is_configured(cls):
145154
return cls.current.settings is not None

Diff for: taskbadger/sdk.py

+35-21
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
from taskbadger.exceptions import (
66
ConfigurationError,
7+
MissingConfiguration,
78
ServerError,
89
TaskbadgerException,
910
Unauthorized,
@@ -21,11 +22,11 @@
2122
PatchedTaskRequestTags,
2223
StatusEnum,
2324
TaskRequest,
24-
TaskRequestTags,
2525
)
2626
from taskbadger.internal.types import UNSET
27-
from taskbadger.mug import Badger, Session, Settings
27+
from taskbadger.mug import Badger, Callback, Session, Settings
2828
from taskbadger.systems import System
29+
from taskbadger.utils import import_string
2930

3031
log = 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

4951
def _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

Diff for: taskbadger/utils.py

+15
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
from importlib import import_module
2+
3+
4+
def import_string(dotted_path):
5+
try:
6+
module_path, class_name = dotted_path.rsplit(".", 1)
7+
except ValueError as err:
8+
raise ImportError("%s doesn't look like a module path" % dotted_path) from err
9+
10+
module = import_module(module_path)
11+
12+
try:
13+
return getattr(module, class_name)
14+
except AttributeError as err:
15+
raise ImportError(f'Module "{module_path}" does not define a "{class_name}" attribute/class') from err

Diff for: tests/test_init.py

+30
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
import pytest
2+
3+
from taskbadger import Badger, init
4+
from taskbadger.exceptions import ConfigurationError
5+
from taskbadger.mug import _local
6+
7+
8+
@pytest.fixture(autouse=True)
9+
def _reset():
10+
b_global = Badger.current
11+
_local.set(Badger())
12+
yield
13+
_local.set(b_global)
14+
15+
16+
def test_init():
17+
init("org", "project", "token", before_create=lambda x: x)
18+
19+
20+
def test_init_import_before_create():
21+
init("org", "project", "token", before_create="tests.test_init._before_create")
22+
23+
24+
def test_init_import_before_create_fail():
25+
with pytest.raises(ConfigurationError):
26+
init("org", "project", "token", before_create="missing")
27+
28+
29+
def _before_create(_):
30+
pass

Diff for: tests/test_sdk.py

+37-1
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
import pytest
55

6-
from taskbadger import Action, EmailIntegration, StatusEnum, WebhookIntegration
6+
from taskbadger import Action, EmailIntegration, StatusEnum, WebhookIntegration, create_task
77
from taskbadger.exceptions import TaskbadgerException
88
from taskbadger.internal.models import (
99
PatchedTaskRequest,
@@ -95,6 +95,42 @@ def test_create(settings, patched_create):
9595
)
9696

9797

98+
def test_before_create_update_task(settings, patched_create):
99+
def before_create(task):
100+
tags = task.setdefault("tags", {})
101+
tags["new"] = "tag"
102+
return task
103+
104+
settings.before_create = before_create
105+
106+
api_task = task_for_test()
107+
patched_create.return_value = Response(HTTPStatus.OK, b"", {}, api_task)
108+
109+
task = create_task(name="task name")
110+
assert task.id == api_task.id
111+
112+
request = TaskRequest.from_dict(
113+
{
114+
"name": "task name",
115+
"status": StatusEnum.PENDING,
116+
"tags": {"new": "tag"},
117+
}
118+
)
119+
assert patched_create.call_args[1]["body"] == request
120+
121+
122+
def test_before_create_filter(settings, patched_create):
123+
def before_create(_):
124+
return None
125+
126+
settings.before_create = before_create
127+
128+
with pytest.raises(TaskbadgerException):
129+
create_task(name="task name")
130+
131+
patched_create.assert_not_called()
132+
133+
98134
def test_update_status(settings, patched_update):
99135
api_task = task_for_test()
100136
task = Task(api_task)

0 commit comments

Comments
 (0)