[Python Context and ContextVars] #python #python3 #context #contextvars

0. README.md

Context variables are variables that can have different values depending on their context. They are similar to Thread-Local Storage in which each execution thread may have a different value for a variable. However, with context variables, there may be several contexts in one execution thread. The main use case for context variables is keeping track of variables in concurrent asynchronous tasks. -- https://realpython.com/python37-new-features/#context-variables

"""Example copied verbatim from Real Python."""

import contextvars

name = contextvars.ContextVar("name")
contexts = list()

def greet():
    print(f"Hello {name.get()}")

# Construct contexts and set the context variable name
for first_name in ["Steve", "Dina", "Harry"]:
    ctx = contextvars.copy_context()
    ctx.run(name.set, first_name)
    contexts.append(ctx)

# Run greet function inside each context
for ctx in reversed(contexts):
    ctx.run(greet)

1. Python ContextVars.py

import contextvars

"""
Important: Context Variables should be created at the top module level and never in closures. 
Context objects (we'll see in the next file) hold strong references to context variables.
Scoped ContextVars prevents those context variables from being properly garbage collected.
"""

var = contextvars.ContextVar("foo")

var.get("foo")  # 'foo' (no value is set, so we just get the 'name' of the variable back)

"""
NOTE:

Naming the variable `var` is actually a bit confusing/misleading. 
It should really be named after the value it will contain.
A more practical example would be `id = contextvars.ContextVar("id")`.
Then you would do `id.set("123")`

But for the sake of testing this code in a REPL, I opted for just naming it `var` instead.
"""

token = var.set("bar")
token.old_value  # <Token.MISSING>

var.get("foo")  # 'bar'

token2 = var.set("baz")
token.old_value  # 'bar'

var.get("foo")  # 'baz'

var.reset(token2)
var.get("foo")  # 'bar'

var.reset(token)
var.get("foo")  # 'foo' (i.e. no value)

"""
NOTE:

I could have reset `var` in a different order.
I didn't have to reset using `token2` then `token`.
I could have reset using `token` first, then `token2`.
Doing that would have meant `var` would still have a value set of 'bar' (as per `token2.old_value`)

The following code presumes the latter was done (i.e. `token2` was used as the last `var.reset()` token)
"""

2. Python Context.py

"""
In the following code we look at the contextvars.Context object, which is a mapping of ContextVars to their values.

Whenever you import the contextvars module you'll find that there is 'default' Context created.

If you set a ContextVar in any modules that have imported the contextvars module, then you'll discover the
default Context is shared between modules and so it'll show the same ContextVar across all modules.

You can access the default Context by taking a copy of it (see below).

It's important to realize that defining a ContextVar will not mean it shows up in the Context _unless_ 
you set a value onto the ContextVar. Because the following code presumes the earlier code in file 1. 
was executed, it means we can see the 'foo' ContextVar that was set.
"""

ctx = contextvars.copy_context()  # <Context at 0x106e23840>

list(ctx.keys())  # [<ContextVar name='foo' at 0x106f52590>]
list(ctx.items())  # [(<ContextVar name='foo' at 0x106f52590>, 'bar')]

"""
Context() creates an empty context with no values in it.
"""

newctx = contextvars.Context()  # <Context at 0x106fa2ac0>

list(newctx.items())  # []

"""
Changes can be made to a Context's ContextVar(s) if modified via the contextvars.Context().run() method

The following code snippet presumes a fresh environment (no previous Context or ContextVars)...
"""

var = contextvars.ContextVar('foo')
var.set('bar')

def scope():
    var.set('baz')
    print(var.get('foo'))  # 'baz'
    print(ctx.get(var))  # 'baz'
    return "finished"

ctx = contextvars.copy_context()

list(ctx.items())  # [(<ContextVar name='foo' at 0x1025a1450>, 'bar')]

result = ctx.run(scope)  # 'finished'

"""
NOTE:

If you're just doing a WRITE operation then pass the `.set()` method to `.run()`

e.g. ctx.run(var.set, 'baz')
"""

list(ctx.items())  # [(<ContextVar name='foo' at 0x1025a1450>, 'baz')]

var.get('foo')  # 'bar'

"""
Unforunately the object model is a bit crappy and so it's not easy to get at the internal ContextVars a Context holds.

I wrote a quick lookup function to help with that...
"""

import contextvars
from typing import Optional


def lookup(ctx: contextvars.Context, key: str) -> Optional[str]:
    for i, v in list(ctx.items()):
        if i.name == key:
            return v
    return None
  
lookup(ctx, "foo")  # 'bar'

3. Server Example.py

"""
I wanted to try and mimick something like golang's context.Context 
which is built-in to their http server by default.

I'm sort of surprised Python hasn't tried to copy that approach?

We've got three files in this example...

1. ctx.py: abstraction for contextvars module
2. foo.py: random module for generating an ID
3. app.py: web server module using asyncio
"""

# ctx.py
#
import contextvars
from typing import Optional


def lookup(ctx: contextvars.Context, key: str) -> Optional[str]:
    for i, v in list(ctx.items()):
        if i.name == key:
            return v
    return None


def new() -> contextvars.Context:
    return contextvars.Context()


# foo.py
#
import asyncio
import contextvars
import os
import random

id: contextvars.ContextVar = contextvars.ContextVar('id')


def gen_id():
    uid = str(os.urandom(15))
    print("uid:", uid)
    id.set(uid)


async def bar(ctx: contextvars.Context):
    ctx.run(gen_id)
    r = random.randint(5, 10)
    print(f"sleep for: {r} seconds")
    await asyncio.sleep(r)


# app.py
#
import asyncio

import ctx
import foo


async def handle_request(reader, writer):
    c = ctx.new()
    await foo.bar(c)
    writer.write(f"result: {ctx.lookup(c, 'id')}".encode())
    writer.close()


async def main():
    srv = await asyncio.start_server(handle_request, '127.0.0.1', 8081)

    async with srv:
        await srv.serve_forever()

asyncio.run(main())

Unique context vars are created per object and not per name. If some-other module creates context var with same name inside your context, it'd mess with your lookup.

https://gist.github.com/Integralist/4e2f323f29ceb624f7fd540687d8e74f#file-2-python-context-py-L62-L78

import contextvars


foo = contextvars.ContextVar('var')
bar = contextvars.ContextVar('var')
ctx_map = {
    'foo': foo,
    'bar': bar
}

if __name__ == '__main__':
    print(foo)  # <ContextVar name='var' at 0x10d3ea2f0>
    print(bar)  # <ContextVar name='var' at 0x10d3ea230>
    print(ctx_map['foo'])  # <ContextVar name='var' at 0x10d3ea2f0>
    print(ctx_map['bar'])  # <ContextVar name='var' at 0x10d3ea230>
    print(foo.get('None'))  # None
    print(bar.get('None'))  # None
    print(ctx_map['foo'].get('None'))  # None
    print(ctx_map['bar'].get('None'))  # None

    foo.set('Foo')
    bar.set('Bar')
    print(foo)  # <ContextVar name='var' at 0x10d3ea2f0>
    print(bar)  # <ContextVar name='var' at 0x10d3ea230>
    print(ctx_map['foo'])  # <ContextVar name='var' at 0x10d3ea2f0>
    print(ctx_map['bar'])  # <ContextVar name='var' at 0x10d3ea230>
    print(foo.get('None'))  # Foo
    print(bar.get('None'))  # Bar
    print(ctx_map['foo'].get('None'))  # Foo
    print(ctx_map['bar'].get('None'))  # Bar

    ctx_map['foo'].set('FooFoo')
    ctx_map['bar'].set('BarBar')
    print(foo)  # <ContextVar name='var' at 0x10d3ea2f0>
    print(bar)  # <ContextVar name='var' at 0x10d3ea230>
    print(ctx_map['foo'])  # <ContextVar name='var' at 0x10d3ea2f0>
    print(ctx_map['bar'])  # <ContextVar name='var' at 0x10d3ea230>
    print(foo.get('None'))  # FooFoo
    print(bar.get('None'))  # BarBar
    print(ctx_map['foo'].get('None'))  # FooFoo
    print(ctx_map['bar'].get('None'))  # BarBar