Python contextvars.ContextVar() Examples

def coroutine(self, coroutine):
        create_call_trace.append('SpecialCoroutineMiddleware')

        # noinspection PyCompatibility
        async def wrapper():
            var = contextvars.ContextVar('middleware_var')  # type: contextvars.ContextVar[int]
            var.set(self.var_value)

            run_call_trace_pre.append('SpecialCoroutineMiddleware')

            try:
                return await coroutine
            finally:
                run_call_trace_post.append('SpecialCoroutineMiddleware')

        return wrapper() 

def test_context_threads_1(self):
        cvar = contextvars.ContextVar('cvar')

        def sub(num):
            for i in range(10):
                cvar.set(num + i)
                time.sleep(random.uniform(0.001, 0.05))
                self.assertEqual(cvar.get(), num + i)
            return num

        tp = concurrent.futures.ThreadPoolExecutor(max_workers=10)
        try:
            results = list(tp.map(sub, range(10)))
        finally:
            tp.shutdown()
        self.assertEqual(results, list(range(10)))


# HAMT Tests 

def test_context_copy_1(self):
        ctx1 = contextvars.Context()
        c = contextvars.ContextVar('c', default=42)

        def ctx1_fun():
            c.set(10)

            ctx2 = ctx1.copy()
            self.assertEqual(ctx2[c], 10)

            c.set(20)
            self.assertEqual(ctx1[c], 20)
            self.assertEqual(ctx2[c], 10)

            ctx2.run(ctx2_fun)
            self.assertEqual(ctx1[c], 20)
            self.assertEqual(ctx2[c], 30)

        def ctx2_fun():
            self.assertEqual(c.get(), 10)
            c.set(30)
            self.assertEqual(c.get(), 30)

        ctx1.run(ctx1_fun) 

def test_context_getset_3(self):
        c = contextvars.ContextVar('c', default=42)
        ctx = contextvars.Context()

        def fun():
            self.assertEqual(c.get(), 42)
            with self.assertRaises(KeyError):
                ctx[c]
            self.assertIsNone(ctx.get(c))
            self.assertEqual(ctx.get(c, 'spam'), 'spam')
            self.assertNotIn(c, ctx)
            self.assertEqual(list(ctx.keys()), [])

            t = c.set(1)
            self.assertEqual(list(ctx.keys()), [c])
            self.assertEqual(ctx[c], 1)

            c.reset(t)
            self.assertEqual(list(ctx.keys()), [])
            with self.assertRaises(KeyError):
                ctx[c]

        ctx.run(fun) 

def test_context_var_repr_1(self):
        c = contextvars.ContextVar('a')
        self.assertIn('a', repr(c))

        c = contextvars.ContextVar('a', default=123)
        self.assertIn('123', repr(c))

        lst = []
        c = contextvars.ContextVar('a', default=lst)
        lst.append(c)
        self.assertIn('...', repr(c))
        self.assertIn('...', repr(lst))

        t = c.set(1)
        self.assertIn(repr(c), repr(t))
        self.assertNotIn(' used ', repr(t))
        c.reset(t)
        self.assertIn(' used ', repr(t)) 

def reset(self, token):  # type: (Token) -> None
        if token.context_var is not self:
            raise ValueError('Token was created by a different ContextVar')
        if self._cv_variable is not None:
            if not isinstance(token, _ContextVarToken):
                raise TypeError('Unexpected `{}` expecting `_ContextVarToken`'.format(token.__class__.__name__))
            self._cv_variable.reset(token.token)
        elif self._tl_variable is not None:
            if not isinstance(token, _ThreadLocalToken):
                raise TypeError('Unexpected `{}` expecting `_ThreadLocalToken`'.format(token.__class__.__name__))
            if token.previous_value is _NO_DEFAULT:
                delattr(self._tl_variable, 'value')
            else:
                self._tl_variable.value = token.previous_value
        else:
            raise TypeError('This context var has been internally messed with and is no longer valid.') 

def _make_threadlocal_contextvars(local):
    # type: (type) -> type
    class ContextVar(object):
        # Super-limited impl of ContextVar

        def __init__(self, name):
            # type: (str) -> None
            self._name = name
            self._local = local()

        def get(self, default):
            # type: (Any) -> Any
            return getattr(self._local, "value", default)

        def set(self, value):
            # type: (Any) -> None
            self._local.value = value

    return ContextVar 

def test_context_var_new_1(self):
        with self.assertRaisesRegex(TypeError, 'takes exactly 1'):
            contextvars.ContextVar()

        with self.assertRaisesRegex(TypeError, 'must be a str'):
            contextvars.ContextVar(1)

        c = contextvars.ContextVar('aaa')
        self.assertEqual(c.name, 'aaa')

        with self.assertRaises(AttributeError):
            c.name = 'bbb'

        self.assertNotEqual(hash(c), hash('aaa')) 

def test_contextvars(self):
        from contextvars import ContextVar
        var = ContextVar('var')
        var.set(0)

        async def set_val():
            var.set(42)

        async def coro():
            await set_val()
            await asyncio.sleep(0.01)
            return var.get()

        result = self.loop.run_until_complete(coro())
        self.assertEqual(result, 42) 

def _is_contextvars_broken():
    # type: () -> bool
    """
    Returns whether gevent/eventlet have patched the stdlib in a way where thread locals are now more "correct" than contextvars.
    """
    try:
        from gevent.monkey import is_object_patched  # type: ignore

        if is_object_patched("threading", "local"):
            # Gevent 20.5 is able to patch both thread locals and contextvars,
            # in that case all is good.
            if is_object_patched("contextvars", "ContextVar"):
                return False

            return True
    except ImportError:
        pass

    try:
        from eventlet.patcher import is_monkey_patched  # type: ignore

        if is_monkey_patched("thread"):
            return True
    except ImportError:
        pass

    return False 

def test_context_getset_5(self):
        c = contextvars.ContextVar('c', default=42)
        c.set([])

        def fun():
            c.set([])
            c.get().append(42)
            self.assertEqual(c.get(), [42])

        contextvars.copy_context().run(fun)
        self.assertEqual(c.get(), []) 

def test_context_getset_2(self):
        v1 = contextvars.ContextVar('v1')
        v2 = contextvars.ContextVar('v2')

        t1 = v1.set(42)
        with self.assertRaisesRegex(ValueError, 'by a different'):
            v2.reset(t1) 

def test_context_run_6(self):
        ctx = contextvars.Context()
        c = contextvars.ContextVar('a', default=0)

        def fun():
            self.assertEqual(c.get(), 0)
            self.assertIsNone(ctx.get(c))

            c.set(42)
            self.assertEqual(c.get(), 42)
            self.assertEqual(ctx.get(c), 42)

        ctx.run(fun) 

def test_context_run_5(self):
        ctx = contextvars.Context()
        var = contextvars.ContextVar('var')

        def func():
            self.assertIsNone(var.get(None))
            var.set('spam')
            1 / 0

        with self.assertRaises(ZeroDivisionError):
            ctx.run(func)

        self.assertIsNone(var.get(None)) 

def test_context_typerrors_1(self):
        ctx = contextvars.Context()

        with self.assertRaisesRegex(TypeError, 'ContextVar key was expected'):
            ctx[1]
        with self.assertRaisesRegex(TypeError, 'ContextVar key was expected'):
            1 in ctx
        with self.assertRaisesRegex(TypeError, 'ContextVar key was expected'):
            ctx.get(1) 

def test_context_subclassing_1(self):
        with self.assertRaisesRegex(TypeError, 'not an acceptable base type'):
            class MyContextVar(contextvars.ContextVar):
                # Potentially we might want ContextVars to be subclassable.
                pass

        with self.assertRaisesRegex(TypeError, 'not an acceptable base type'):
            class MyContext(contextvars.Context):
                pass

        with self.assertRaisesRegex(TypeError, 'not an acceptable base type'):
            class MyToken(contextvars.Token):
                pass 

def _get_contextvars():
    # type: () -> Tuple[bool, type]
    """
    Figure out the "right" contextvars installation to use. Returns a
    `contextvars.ContextVar`-like class with a limited API.

    See https://docs.sentry.io/platforms/python/contextvars/ for more information.
    """
    if not _is_contextvars_broken():
        # aiocontextvars is a PyPI package that ensures that the contextvars
        # backport (also a PyPI package) works with asyncio under Python 3.6
        #
        # Import it if available.
        if sys.version_info < (3, 7):
            # `aiocontextvars` is absolutely required for functional
            # contextvars on Python 3.6.
            try:
                from aiocontextvars import ContextVar  # noqa

                return True, ContextVar
            except ImportError:
                pass
        else:
            # On Python 3.7 contextvars are functional.
            try:
                from contextvars import ContextVar

                return True, ContextVar
            except ImportError:
                pass

    # Fall back to basic thread-local usage.

    from threading import local

    return False, _make_threadlocal_contextvars(local) 

def test_context_var_new_2(self):
        self.assertIsNone(contextvars.ContextVar[int]) 

def __init__(self) -> None:
        self._current_context = ContextVar(
            self._CONTEXT_KEY, default=Context()
        ) 

def __init__(self, peer: Peer):
            self._cached_funcs = {}
            self.peer = peer
            self.order_key = ContextVar('order_key', default=None) 

def __init__(self, context_var, token):  # type: (ContextVar, contextvars.Token) -> None
        super(_ContextVarToken, self).__init__(context_var)
        self.token = token 

def __init__(self, context_var):  # type: (ContextVar) -> None
        self.context_var = context_var 

def __repr__(self):  # type: () -> six.text_type
        r = "<ContextVar name='{}'".format(self.name)
        if self._default is not _NO_DEFAULT:
            r += ' default={!r}'.format(self._default)
        return r + ' at {:0x}>'.format(id(self)) 

def __init__(self, name):
        """Initialize a StackVar."""
        self.var = ContextVar(name, default=(None, None))
        self.var.set((None, None)) 

def __init__(self, name, default=cast(VT, _NO_DEFAULT)):  # type: (six.text_type, VT) -> None
        self.name = name
        self._has_default = default is not _NO_DEFAULT
        self._default = default

        self._cv_variable = None  # type: Optional[contextvars.ContextVar[VT]]
        self._tl_variable = None  # type: Optional[threading.local]
        if contextvars:
            if self._has_default:
                self._cv_variable = contextvars.ContextVar(name, default=default)
            else:
                self._cv_variable = contextvars.ContextVar(name)
        elif threading:
            self._tl_variable = threading.local()
        # else is not possible 

def context(
        values: Iterable[Tuple[contextvars.ContextVar[Any], Any]],
) -> Iterator[None]:
    """
    A context manager to set the context variables temporarily.
    """
    tokens: List[Tuple[contextvars.ContextVar[Any], contextvars.Token[Any]]] = []
    try:
        for var, val in values:
            token = var.set(val)
            tokens.append((var, token))
        yield
    finally:
        for var, token in reversed(tokens):
            var.reset(token) 

def __init__(self, name, default):
            self.name = name
            self.contextvar = contextvars.ContextVar(name)
            self.default = default if callable(default) else (lambda: default) 

def __init_subclass__(cls, **kwargs):
        cls.__context_instance = contextvars.ContextVar(
            "instance_" + (kwargs.get("ctx_name") or cls.__name__)
        )
        return cls 

def test_one_thread():
    var = ContextVar('test_one_thread1')  # type: ContextVar[six.text_type]
    assert "<ContextVar name='test_one_thread1' at " in repr(var)
    assert 'default=' not in repr(var)

    with pytest.raises(LookupError) as error_context:
        var.get()

    if six.PY2:
        assert error_context.value.args[0] is var
        assert var._tl_variable is not None
        assert isinstance(var._tl_variable, getattr(threading, 'local'))
    else:
        assert var._cv_variable is not None
        assert contextvars is not None
        assert isinstance(var._cv_variable, contextvars.ContextVar)

    assert var.get('default1') == 'default1'
    assert var.get(default='default2') == 'default2'

    var.set('set1')
    assert var.get() == 'set1'

    var = ContextVar('test_one_thread2', 'default3')
    if six.PY2:
        assert "<ContextVar name='test_one_thread2' default=u'default3' at " in repr(var)
    else:
        assert "<ContextVar name='test_one_thread2' default='default3' at " in repr(var)

    assert var.get() == 'default3'
    assert var.get('default4') == 'default4'

    var.set('set2')
    assert var.get() == 'set2'

    var = ContextVar('test_one_thread3', default='default5')

    assert var.get() == 'default5'
    assert var.get('default6') == 'default6'

    var.set('set3')
    assert var.get() == 'set3' 

def test_coroutine_middleware():
    before_call_trace.clear()
    create_call_trace.clear()
    run_call_trace_pre.clear()
    run_call_trace_post.clear()

    var = contextvars.ContextVar('caller_var')  # type: contextvars.ContextVar[str]
    var.set('yes man')

    test_context = {
        'caller_var': None,
        'middleware_var': None
    }  # type: Dict[str, Any]

    # noinspection PyCompatibility
    async def coroutine():
        run_call_trace_pre.append('target')

        test_context['caller_var'] = var.get('default_cv')
        for context_var in contextvars.copy_context().keys():
            if context_var.name == 'middleware_var':
                test_context['middleware_var'] = context_var.get('default_mv')
        await asyncio.sleep(0.05)

        run_call_trace_post.append('target')

        return 'foo_coroutine_returned_this'

    thread = AsyncEventLoopThread([
        SpecialCoroutineMiddleware(42),
        TracingCoroutineMiddleware(),
    ])
    thread.start()

    future = thread.run_coroutine(coroutine())

    await asyncio.sleep(0.01)

    assert future.result() == 'foo_coroutine_returned_this'

    thread.join()

    assert before_call_trace == ['SpecialCoroutineMiddleware', 'TracingCoroutineMiddleware']
    assert create_call_trace == ['TracingCoroutineMiddleware', 'SpecialCoroutineMiddleware']
    assert run_call_trace_pre == ['SpecialCoroutineMiddleware', 'TracingCoroutineMiddleware', 'target']
    assert run_call_trace_post == ['target', 'TracingCoroutineMiddleware', 'SpecialCoroutineMiddleware']

    assert test_context['caller_var'] == 'yes man'
    assert test_context['middleware_var'] == 42


# noinspection PyCompatibility