diff --git a/taskbadger/__init__.py b/taskbadger/__init__.py index 456570c..a03fede 100644 --- a/taskbadger/__init__.py +++ b/taskbadger/__init__.py @@ -1,7 +1,7 @@ from .decorators import track from .integrations import Action, EmailIntegration, WebhookIntegration from .internal.models import StatusEnum -from .mug import Session +from .mug import Badger, Session from .safe_sdk import create_task_safe, update_task_safe from .sdk import DefaultMergeStrategy, Task, create_task, get_task, init, update_task @@ -15,3 +15,7 @@ __version__ = importlib_metadata.version(__name__) except importlib_metadata.PackageNotFoundError: __version__ = "dev" + + +def current_scope(): + return Badger.current.scope() diff --git a/taskbadger/mug.py b/taskbadger/mug.py index d965016..b610d5d 100644 --- a/taskbadger/mug.py +++ b/taskbadger/mug.py @@ -73,6 +73,29 @@ def __exit__(self, *args, **kwargs): self.client = None +class Scope: + """Scope holds global data which will be added to every task created within the current scope. + + Scope data will be merged with task data when creating a task where data provided directly to the task + will override scope data. + """ + + def __init__(self): + self.stack = [] + self.context = {} + + def __enter__(self): + self.stack.append(self.context) + self.context = self.context.copy() + return self + + def __exit__(self, *args): + self.context = self.stack.pop() + + def __setitem__(self, key, value): + self.context[key] = value + + class MugMeta(type): @property def current(cls): @@ -94,6 +117,7 @@ def __init__(self, settings_or_mug=None): self.settings = settings_or_mug self._session = ReentrantSession() + self._scope = Scope() def bind(self, settings): self.settings = settings @@ -104,6 +128,9 @@ def session(self) -> ReentrantSession: def client(self) -> AuthenticatedClient: return self.settings.get_client() + def scope(self) -> Scope: + return self._scope + @classmethod def is_configured(cls): return cls.current.settings is not None diff --git a/taskbadger/sdk.py b/taskbadger/sdk.py index e264651..b051e7d 100644 --- a/taskbadger/sdk.py +++ b/taskbadger/sdk.py @@ -100,8 +100,10 @@ def create_task( task = TaskRequest( name=name, status=status, value=value, value_max=value_max, max_runtime=max_runtime, stale_timeout=stale_timeout ) - if data: - task.data = TaskRequestData.from_dict(data) + scope_data = Badger.current.scope().context + if scope_data or data: + data = data or {} + task.data = TaskRequestData.from_dict({**scope_data, **data}) if actions: task.additional_properties = {"actions": [a.to_dict() for a in actions]} kwargs = _make_args(json_body=task) diff --git a/tests/test_scope.py b/tests/test_scope.py new file mode 100644 index 0000000..8e67c9b --- /dev/null +++ b/tests/test_scope.py @@ -0,0 +1,58 @@ +import random +import threading +import time + +import pytest + +from taskbadger import create_task, init +from taskbadger.mug import GLOBAL_MUG, Badger +from tests.test_sdk_primatives import _json_task_response, _verify_task + + +def test_scope_singleton(): + assert Badger.current == GLOBAL_MUG + scope = Badger.current.scope() + assert scope.context == {} + assert scope.stack == [] + assert scope == Badger.current.scope() + + +def test_scope_context(): + scope = Badger.current.scope() + assert scope.context == {} + assert scope.stack == [] + with scope: + assert scope.stack == [{}] + scope.context["foo"] = "bar" + with scope: + assert scope.stack == [{}, {"foo": "bar"}] + assert scope.context == {"foo": "bar"} + scope.context["bar"] = "bazz" + with scope: + assert scope.context == {"foo": "bar", "bar": "bazz"} + scope.context.clear() + assert scope.context == {"foo": "bar"} + assert scope.stack == [{}] + assert scope.context == {} + assert scope.stack == [] + + +@pytest.fixture(autouse=True) +def init_skd(): + init("org", "project", "token") + + +def test_create_task_with_scope(httpx_mock): + with Badger.current.scope() as scope: + scope["foo"] = "bar" + scope["bar"] = "bazz" + httpx_mock.add_response( + url="https://taskbadger.net/api/org/project/tasks/", + method="POST", + match_headers={"Authorization": "Bearer token"}, + match_content=b'{"name": "name", "status": "pending", "data": {"foo": "bar", "bar": "buzzer"}}', + json=_json_task_response(), + status_code=201, + ) + task = create_task("name", data={"bar": "buzzer"}) + _verify_task(task)