python如何使用contextvars模块源码分析

from flask import Flask, request
app = Flask(__name__)
@app.route('/')
def root():
    so1n_name = request.get('so1n_name')
    return f'Name is {so1n_name}'

from django.http import HttpResponse
def root(request):
    so1n_name = request.get('so1n_name')
    return HttpResponse(f'Name is {so1n_name}')

# 伪代码,举个例子一个request传了3个函数
from django.http import HttpResponse
def is_allow(request, uid):
    if request.ip == '127.0.0.1' and check_permissions(uid):
        return True
    else:
        return False
def check_permissions(request, uid):
    pass

def root(request):
    user_id = request.GET.get('uid')
    if is_allow(request, id):
    	return HttpResponse('ok')
    else
        return HttpResponse('error')

import logging
from typing import Any
from flask import g  # type: ignore
from flask.logging import default_handler
# 这是一个Python logging.Filter的对象, 日志在生成之前会经过Filter步骤, 这时候我们可以为他绑定request_id变量
class RequestIDLogFilter(logging.Filter):
    """
    Log filter to inject the current request id of the request under `log_record.request_id`
    """
    def filter(self, record: Any) -> Any:
        record.request_id = g.request_id or None
        return record
# 配置日志的format格式, 这里多配了一个request_id变量
format_string: str = (
    "[%(asctime)s][%(levelname)s][%(filename)s:%(lineno)d:%(funcName)s:%(request_id)s]" " %(message)s"
)
# 为flask的默认logger设置format和增加一个logging.Filter对象
default_handler.setFormatter(logging.Formatter(format_string))
default_handler.addFilter(RequestIDLogFilter())
# 该方法用于设置request_id
def set_request_id() -> None:
    g.request_id = request.headers.get("X-Request-Id", str(uuid4()))
# 初始化FLask对象, 并设置before_request
app: Flask = Flask("demo")
app.before_request(set_request_id)

# demo/web_tools.py
# 通过中间件把变量给存进去
class RequestContextMiddleware(BaseHTTPMiddleware):
    async def dispatch(
            self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        request.stat.redis = REDIS_POOL
        response = await call_next(request)
        return response
# demo/server.py
# 调用变量
@APP.route('/')
async def homepage(request):
    # 伪代码,这里是执行redis命令
    await request.stat.redis.execute()
    return JSONResponse({'hello': 'world'})

# demo/context.py
# 该文件存放contextvars相关
import contextvars
if TYPE_CHECKING:
    from demo.redis_dal import RDS  # 这里是一个redis的封装实例
# 初始化一个redis相关的全局context
redis_pool_context = contextvars.ContextVar('redis_pool')
# 通过函数调用可以获取到当前协程运行时的context上下文
def get_redis() -> 'RDS':
    return redis_pool_context.get()
# demo/web_tool.py
# 该文件存放starlette相关模块
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.middleware.base import RequestResponseEndpoint
from starlette.responses import Response
from demo.redis_dal import RDS
# 初始化一个redis客户端变量,当前为空
REDIS_POOL = None  # type: Optional[RDS]
class RequestContextMiddleware(BaseHTTPMiddleware):
    async def dispatch(
            self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        # 通过中间件,在进入路由之前,把redis客户端放入当前协程的上下文之中
        token = redis_pool_context.set(REDIS_POOL)
        try:
        	response = await call_next(request)
            return response
        finally:
        	# 调用完成,回收当前请求设置的redis客户端的上下文
            redis_pool_context.reset(token)
async def startup_event() -> None:
    global REDIS_POOL
    REDIS_POOL = RDS() # 初始化客户端,里面通过asyncio.ensure_future逻辑延后连接
async def shutdown_event() -> None:
    if REDIS_POOL:
        await REDIS_POOL.close() # 关闭redis客户端
# demo/server.py
# 该文件存放starlette main逻辑
from starlette.applications import Starlette
from starlette.responses import JSONResponse
from demo.web_tool import RequestContextMiddleware
from demo.context import get_redis
APP = Starlette()
APP.add_middleware(RequestContextMiddleware)
@APP.route('/')
async def homepage(request):
    # 伪代码,这里是执行redis命令
    # 只要验证 id(get_redis())等于demo.web_tool里REDID_POOL的id一致,那证明contextvars可以为asyncio维护一套上下文状态
    await get_redis().execute()
    return JSONResponse({'hello': 'world'})

import asyncio
import uuid
from contextvars import Context, copy_context
from functools import partial
from typing import Optional, Set
import httpx
from fastapi import FastAPI, Request, Response
from fast_tools.context import ContextBaseModel, ContextMiddleware, HeaderHelper
app: FastAPI = FastAPI()
check_set: Set[int] = set()
class ContextModel(ContextBaseModel):
    """
	通过该实例可以屏蔽大部分与contextvars相关的操作,如果要添加一个变量,则在该实例添加一个属性即可.
	属性必须要使用Type Hints的写法,不然不会识别(强制使用Type Hints)
    """
    # 用于把自己的实例(如上文所说的redis客户端)存放于contextvars中
    http_client: httpx.AsyncClient
    # HeaderHepler用于把header的变量存放于contextvars中
    request_id: str = HeaderHelper.i("X-Request-Id", default_func=lambda request: str(uuid.uuid4()))
    ip: str = HeaderHelper.i("X-Real-IP", default_func=lambda request: request.client.host)
    user_agent: str = HeaderHelper.i("User-Agent")

    async def before_request(self, request: Request) -> None:
        # 请求之前的钩子, 通过该钩子可以设置自己的变量
        self.http_client = httpx.AsyncClient()
        check_set.add(id(self.http_client))

    async def before_reset_context(self, request: Request, response: Optional[Response]) -> None:
        # 准备退出中间件的钩子, 这步奏后会清掉上下文
        await self.http_client.aclose()
context_model: ContextModel = ContextModel()
app.add_middleware(ContextMiddleware, context_model=context_model)
async def test_ensure_future() -> None:
    assert id(context_model.http_client) in check_set
def test_run_in_executor() -> None:
    assert id(context_model.http_client) in check_set
def test_call_soon() -> None:
    assert id(context_model.http_client) in check_set
@app.get("/")
async def root() -> dict:
    # 在使用asyncio.ensure_future开启另外一个子协程跑任务时, 也可以复用上下文
    asyncio.ensure_future(test_ensure_future())
    loop: "asyncio.AbstractEventLoop" = asyncio.get_event_loop()
    # 使用call_soon也能复用上下文
    loop.call_soon(test_call_soon)
    # 使用run_in_executor也能复用上下文, 但必须使用上下文的run方法, copy_context表示复制当前的上下文
    ctx: Context = copy_context()
    await loop.run_in_executor(None, partial(ctx.run, test_run_in_executor))  # type: ignore
    return {
        "message": context_model.to_dict(is_safe_return=True),  # not return CustomQuery
        "client_id": id(context_model.http_client),
    }
if __name__ == "__main__":
    import uvicorn  # type: ignore
    uvicorn.run(app)

class ContextMeta(type(collections.abc.Mapping)):
    # contextvars.Context is not subclassable.
    def __new__(mcls, names, bases, dct):
        cls = super().__new__(mcls, names, bases, dct)
        if cls.__module__ != 'contextvars' or cls.__name__ != 'Context':
            raise TypeError("type 'Context' is not an acceptable base type")
        return cls

class Token(metaclass=TokenMeta):
    MISSING = object()
    def __init__(self, context, var, old_value):
        # 分别存放上下文变量, 当前set的数据以及上次set的数据
        self._context = context
        self._var = var
        self._old_value = old_value
        self._used = False
    @property
    def var(self):
        return self._var
    @property
    def old_value(self):
        return self._old_value
    def __repr__(self):
        r = '<Token '
        if self._used:
            r += ' used'
        r += ' var={!r} at {:0x}>'.format(self._var, id(self))
        return r

def copy_context():
    return _get_context().copy()
def _get_context():
    ctx = getattr(_state, 'context', None)
    if ctx is None:
        ctx = Context()
        _state.context = ctx
    return ctx
def _set_context(ctx):
    _state.context = ctx
_state = threading.local()

pet_dict = {}
def get_pet(pet_name):
    return pet_dict[pet_name]
def set_pet(pet_name):
    return pet_dict[pet_name]

pet_dict = {}
def get_pet(pet_name, pid):
    return pet_dict[pid][pet_name]
def set_pet(pet_name, pid):
    return pet_dict[pid][pet_name]

import threading
thread_local=threading.local()
def get_pet(pet_name):
    return thread_local[pet_name]
def set_pet(pet_name):
    return thread_local[pet_name]

class Context(collections.abc.Mapping, metaclass=ContextMeta):
    def __init__(self):
        self._data = immutables.Map()
        self._prev_context = None
    def run(self, callable, *args, **kwargs):
        if self._prev_context is not None:
            raise RuntimeError(
                'cannot enter context: {} is already entered'.format(self))
        self._prev_context = _get_context()
        try:
            _set_context(self)
            return callable(*args, **kwargs)
        finally:
            _set_context(self._prev_context)
            self._prev_context = None
    def copy(self):
        new = Context()
        new._data = self._data
        return new

map2 = map.set('a', 10)
print(map, map2)
# will print:
#   <immutables.Map({'a': 1, 'b': 2})>
#   <immutables.Map({'a': 10, 'b': 2})>
map3 = map2.delete('b')
print(map, map2, map3)
# will print:
#   <immutables.Map({'a': 1, 'b': 2})>
#   <immutables.Map({'a': 10, 'b': 2})>
#   <immutables.Map({'a': 10})>

def run(self, callable, *args, **kwargs):
    # 已经存在旧的context,抛出异常,防止多线程循环调用
    if self._prev_context is not None:
        raise RuntimeError(
            'cannot enter context: {} is already entered'.format(self))
    self._prev_context = _get_context()  # 保存当前的context
    try:
        _set_context(self) # 设置新的context
        return callable(*args, **kwargs)  # 执行函数
    finally:
        _set_context(self._prev_context)  # 设置为旧的context
        self._prev_context = None

def set(self, value):
    ctx = _get_context()  # 获取当前上下文对象`Context`
    data = ctx._data
    try:
        old_value = data[self]  # 获取Context旧对象
    except KeyError:
        old_value = Token.MISSING  # 获取不到则填充一个object(全局唯一)
    updated_data = data.set(self, value) # 设置新的值
    ctx._data = updated_data
    return Token(ctx, self, old_value) # 返回带有旧值的token

def get(self, default=_NO_DEFAULT):
    ctx = _get_context()  # 获取当前上下文对象`Context`
    try:
        return ctx[self]  # 返回获取的值
    except KeyError:
        pass
    if default is not _NO_DEFAULT:
        return default    # 返回调用get时设置的值
    if self._default is not _NO_DEFAULT:
        return self._default  # 返回初始化context时设置的默认值
    raise LookupError  # 都没有则会抛错

def reset(self, token):
       if token._used:
       	# 判断token是否已经被使用
           raise RuntimeError("Token has already been used once")
       if token._var is not self:
       	# 判断token是否是当前contextvar返回的
           raise ValueError(
               "Token was created by a different ContextVar")
       if token._context is not _get_context():
       	# 判断token的上下文是否跟contextvar上下文一致
           raise ValueError(
               "Token was created in a different Context")
       ctx = token._context
       if token._old_value is Token.MISSING:
       	# 如果没有旧值则删除该值
           ctx._data = ctx._data.delete(token._var)
       else:
       	# 有旧值则当前contextvar变为旧值
           ctx._data = ctx._data.set(token._var, token._old_value)
       token._used = True  # 设置flag,标记token已经被使用了

PY37 = sys.version_info >= (3, 7)
if PY37:
    def asyncio_current_task(loop=None):
        """Return the current task or None."""
        try:
            return asyncio.current_task(loop)
        except RuntimeError:
            # simulate old behaviour
            return None
else:
    asyncio_current_task = asyncio.Task.current_task

def set(key, value):
    """
    Sets the given value inside Task.context[key]. If the key does not exist it creates it.
    :param key: identifier for accessing the context dict.
    :param value: value to store inside context[key].
    :raises
    """
    current_task = asyncio_current_task()
    if not current_task:
        raise ValueError(NO_LOOP_EXCEPTION_MSG.format(key))

    current_task.context[key] = value

def get(key, default=None):
    """
    Retrieves the value stored in key from the Task.context dict. If key does not exist,
    or there is no event loop running, default will be returned
    :param key: identifier for accessing the context dict.
    :param default: None by default, returned in case key is not found.
    :return: Value stored inside the dict[key].
    """
    current_task = asyncio_current_task()
    if not current_task:
        raise ValueError(NO_LOOP_EXCEPTION_MSG.format(key))
    return current_task.context.get(key, default)

def clear():
    """
    Clear the Task.context.
    :raises ValueError: if no current task.
    """
    current_task = asyncio_current_task()
    if not current_task:
        raise ValueError("No event loop found")
    current_task.context.clear()

def task_factory(loop, coro, copy_context=False, context_factory=None):
    """
    By default returns a task factory that uses a simple dict as the task context,
    but allows context creation and inheritance to be customized via ``context_factory``.
    """
    # 生成context工厂函数
    context_factory = context_factory or partial(
        dict_context_factory, copy_context=copy_context)
    # 创建task, 跟asyncio.ensure_future一样
    task = asyncio.tasks.Task(coro, loop=loop)
    if task._source_traceback:
        del [-1]

    # 获取task的context
    try:
        context = asyncio_current_task(loop=loop).context
    except AttributeError:
        context = None
    # 从context工厂中处理context并赋值在task
    task.context = context_factory(context)
    return task

def dict_context_factory(parent_context=None, copy_context=False):
    """A traditional ``dict`` context to keep things simple"""
    if parent_context is None:
        # initial context
        return {}
    else:
        # inherit context
        new_context = parent_context
        if copy_context:
            new_context = deepcopy(new_context)
        return new_context

def chainmap_context_factory(parent_context=None):
    """
    A ``ChainMap`` context, to avoid copying any data
    and yet preserve strict one-way inheritance
    (just like with dict copying)
    """
    if parent_context is None:
        # initial context
        return ChainMap()
    else:
        # inherit context
        if not isinstance(parent_context, ChainMap):
            # if a dict context was previously used, then convert
            # (without modifying the original dict)
            parent_context = ChainMap(parent_context)
        return parent_context.new_child()