diff --git a/gino/__init__.py b/gino/__init__.py index d79eaa54..cf0831bd 100644 --- a/gino/__init__.py +++ b/gino/__init__.py @@ -1,7 +1,13 @@ from .api import Gino from .pool import GinoPool from .connection import GinoConnection -from .local import get_local, enable_task_local, disable_task_local +from .local import ( + get_local, + enable_task_local, + disable_task_local, + reset_local, + is_local_root, +) from .exceptions import * __version__ = '0.5.4' diff --git a/gino/local.py b/gino/local.py index c81c0ac1..34c70d03 100644 --- a/gino/local.py +++ b/gino/local.py @@ -6,21 +6,67 @@ def get_local(): + """Get local storage from current task. + + Returns a dict if task local is enabled for current loop, or ``None`` + otherwise. Please note, if you create a new task from an existing task, + the new task will by default inherit the same local storage object from + the existing task. Use ``is_local_root` to check the inherit state, and use + ``reset_local`` to spawn tasks without inheriting. See the two functions + for more information. + """ return getattr(asyncio.Task.current_task(), 'task_local', None) def task_factory(loop, coro): - orig = _original_task_factories.get(loop) - loop.set_task_factory(orig) - task = loop.create_task(coro) - loop.set_task_factory(task_factory) + loop.set_task_factory(_original_task_factories.get(loop)) + try: + task = loop.create_task(coro) + if getattr(task, '_source_traceback', None): + del getattr(task, '_source_traceback')[-2:] + + local = get_local() + task.task_local = {} if local is None else local + task.task_local_is_root = local is None + + return task + finally: + loop.set_task_factory(task_factory) + + +def reset_local(coro_or_future=None, *, loop=None): + """Reset local to empty dict within given routine, or current if not given. + + Please note, this works only if the local was inherited from another + routine, or this is a no-op. That means, calling ``reset_local`` a second + time on the same routine won't set its local to empty dict, because its + local was already non-inherited after the first call. + + This method always return a future object. New task will be created and + returned if the parameter is a coroutine or awaitable object. Therefore it + is practical to use this method to spawn detached tasks. + """ + if coro_or_future is None: + coro_or_future = asyncio.Task.current_task() + + if asyncio.isfuture(coro_or_future): + if is_local_root(coro_or_future) is False: + coro_or_future.task_local = {} + coro_or_future.task_local_is_root = True + return coro_or_future + else: + return reset_local(asyncio.ensure_future(coro_or_future, loop=loop)) - if getattr(task, '_source_traceback', None): - del getattr(task, '_source_traceback')[-2:] - task.task_local = get_local() or {} +def is_local_root(task=None): + """Check if local of the given task is inherited or not. - return task + If task is not given, current task is used. Returns ``None`` if the task + has no local, or ``True``/``False`` for the inherit state. + """ + if task is None: + task = asyncio.Task.current_task() + return getattr(task, 'task_local_is_root', None) def enable_task_local(loop=None): diff --git a/gino/pool.py b/gino/pool.py index ea3f077d..21e0c4e8 100644 --- a/gino/pool.py +++ b/gino/pool.py @@ -46,6 +46,10 @@ def metadata(self): def execution_options(self): return self._execution_options + @property + def is_root(self): + return self._root is self + async def get_connection(self): if self._closed: raise InterfaceError('This LazyConnection is returned to pool') diff --git a/tests/test_local.py b/tests/test_local.py new file mode 100644 index 00000000..7779f210 --- /dev/null +++ b/tests/test_local.py @@ -0,0 +1,85 @@ +import asyncio + +import pytest +from gino import get_local, reset_local, is_local_root + +pytestmark = pytest.mark.asyncio + + +async def sub(result): + result.append(get_local().get('val')) + + +# noinspection PyUnusedLocal +async def test_attached(task_local): + get_local()['val'] = 123 + result = [] + await asyncio.ensure_future(sub(result)) + assert result[0] == 123 + + +# noinspection PyUnusedLocal +async def test_detached(task_local): + get_local()['val'] = 123 + result = [] + await reset_local(sub(result)) + assert result[0] is None + + +async def test_reset_disabled(): + result = [] + fut = reset_local(sub(result)) + with pytest.raises(AttributeError): + await fut + assert len(result) == 0 + + +async def sub_set(): + get_local()['val'] = 123 + + +# noinspection PyUnusedLocal +async def test_attached_reverse(task_local): + await asyncio.ensure_future(sub_set()) + assert get_local().get('val') == 123 + + +# noinspection PyUnusedLocal +async def test_detached_reverse(task_local): + await reset_local(sub_set()) + assert get_local().get('val') is None + + +# noinspection PyUnusedLocal +async def test_accept_task(task_local): + await reset_local(asyncio.ensure_future(sub_set())) + assert get_local().get('val') is None + + +async def grand_sub(): + assert not is_local_root() + assert get_local().get('val') == 456 + + +async def sub_reset(): + assert is_local_root() is False + assert get_local().get('val') == 123 + + reset_local() + assert is_local_root() + assert get_local().get('val') is None + + get_local()['val'] = 456 + assert get_local().get('val') == 456 + await asyncio.ensure_future(grand_sub()) + + reset_local() + assert get_local().get('val') == 456 + await asyncio.ensure_future(grand_sub()) + + +# noinspection PyUnusedLocal +async def test_inline_reset(task_local): + get_local()['val'] = 123 + await asyncio.ensure_future(sub_reset()) + get_local()['val'] = 123