diff --git a/tasklocals/__init__.py b/tasklocals/__init__.py index d01e532..1a59d9f 100644 --- a/tasklocals/__init__.py +++ b/tasklocals/__init__.py @@ -4,7 +4,14 @@ from weakref import ref import asyncio -__all__ = ['local'] +__all__ = ['local', 'NoTaskError'] + +_default_local_dct = dict() + + +class NoTaskError(RuntimeError): + pass + class _localimpl: """A class managing task-local dicts""" @@ -26,7 +33,7 @@ def get_dict(self): defined.""" task = asyncio.Task.current_task(loop=self._loop) if task is None: - raise RuntimeError("No task is currently running") + raise NoTaskError("No task is currently running") return self.dicts[id(task)][1] def create_dict(self): @@ -35,7 +42,7 @@ def create_dict(self): key = self.key task = asyncio.Task.current_task(loop=self._loop) if task is None: - raise RuntimeError("No task is currently running") + raise NoTaskError("No task is currently running") idt = id(task) def local_deleted(_, key=key): # When the localimpl is deleted, remove the task attribute. @@ -60,25 +67,32 @@ def task_deleted(_, idt=idt): @contextmanager def _patch(self): impl = object.__getattribute__(self, '_local__impl') + strict = object.__getattribute__(self, '_strict') try: dct = impl.get_dict() except KeyError: dct = impl.create_dict() args, kw = impl.localargs self.__init__(*args, **kw) + except NoTaskError: + if strict: + raise + dct = _default_local_dct object.__setattr__(self, '__dict__', dct) yield + class local: - __slots__ = '_local__impl', '__dict__' + __slots__ = '_local__impl', '__dict__', '_strict' - def __new__(cls, *args, loop=None, **kw): + def __new__(cls, *args, loop=None, strict=True, **kw): if (args or kw) and (cls.__init__ is object.__init__): raise TypeError("Initialization arguments are not supported") self = object.__new__(cls) impl = _localimpl(loop=loop) impl.localargs = (args, kw) object.__setattr__(self, '_local__impl', impl) + object.__setattr__(self, '_strict', strict) return self def __getattribute__(self, name): diff --git a/tests/test_local.py b/tests/test_local.py index 26c8565..7f05390 100644 --- a/tests/test_local.py +++ b/tests/test_local.py @@ -2,7 +2,7 @@ import gc import asyncio import asyncio.test_utils -from tasklocals import local +from tasklocals import local, NoTaskError class LocalTests(unittest.TestCase): @@ -16,8 +16,8 @@ def tearDown(self): self.loop.close() gc.collect() - def test_local(self): - mylocal = local(loop=self.loop) + def test_local(self, strict=True): + mylocal = local(loop=self.loop, strict=strict) log1 = [] log2 = [] @@ -62,11 +62,20 @@ def task2(): # ensure all task local values have been properly cleaned up self.assertEqual(object.__getattribute__(mylocal, '_local__impl').dicts, {}) + def test_local_non_strict_mode(self): + self.test_local(strict=False) + def test_local_outside_of_task(self): mylocal = local(loop=self.loop) try: mylocal.foo = 1 - self.fail("RuntimeError has not been raised when tryint to use local object outside of a Task") - except RuntimeError: + self.fail("NoTaskError has not been raised when tryint to use local object outside of a Task") + except NoTaskError: pass + def test_local_outside_of_task_non_strict_mode(self): + mylocal = local(loop=self.loop, strict=False) + try: + mylocal.foo = 1 + except NoTaskError: + self.fail("NoTaskError has been raised when tryint to use non-strict local object outside of a Task")