starlette_context
    plugins
        __init__.py
        plugin.py
        .
        .
        .
    __init__.py
    ctx.py
    header_keys.py
    middleware.py

from contextvars import ContextVar

_request_scope_context_storage: ContextVar[str] = ContextVar(
    "starlette_context"
)

from starlette_context.ctx import context  # noqa: E402, F401

from collections import UserDict
from typing import Any

from contextvars import copy_context
from libs.fd_context import _request_scope_context_storage


class _Context(UserDict):
    """
    A mapping with dict-like interface.
    It is using request context as a data store.
    Can be used only if context has been created in the middleware.

    If you know Flask, it can be compared to g object.
    """

    def __init__(self, *args: Any, **kwargs: Any):
        # not calling super on purpose
        if args or kwargs:
            raise AttributeError("Can't instantiate with attributes")

    @property
    def data(self) -> dict:
        """
        Dump this to json. Object itself it not serializable.
        """
        try:
            return _request_scope_context_storage.get()
        except LookupError as e:
            raise RuntimeError(
                "You didn't use ContextMiddleware or "
                "you're trying to access `context` object "
                "outside of the request-response cycle."
            ) from e

    def exists(self) -> bool:
        return _request_scope_context_storage in copy_context()

    def copy(self) -> dict:
        """
        Read only context data.
        """
        import copy

        return copy.copy(self.data)


context = _Context()

class HeaderKeys:
    correlation_id = "X-Correlation-ID"
    request_id = "X-Request-ID"
    date = "Date"
    forwarded_for = "X-Forwarded-For"
    user_agent = "User-Agent"

from contextvars import Token
from typing import Optional, Sequence

from starlette.middleware.base import (
    BaseHTTPMiddleware,
    RequestResponseEndpoint,
)
from starlette.middleware import Middleware
from starlette.requests import Request
from starlette.responses import Response

from libs.fd_context import _request_scope_context_storage
from libs.fd_context.plugins import Plugin
from libs.fd_context import plugins


class ContextMiddleware(BaseHTTPMiddleware):
    """
    Middleware that creates empty context for request it's used on.
    If not used, you won't be able to use context object.
    """

    def __init__(
        self, plugins: Optional[Sequence[Plugin]] = None, *args, **kwargs
    ) -> None:
        super().__init__(*args, **kwargs)
        self.plugins = plugins or ()
        if not all([isinstance(plugin, Plugin) for plugin in self.plugins]):
            raise TypeError("This is not a valid instance of a plugin")

    async def set_context(self, request: Request) -> dict:
        """
        You might want to override this method.
        The dict it returns will be saved in the scope of a context.
        You can always do that later.
        """
        return {
            plugin.key: await plugin.process_request(request)
            for plugin in self.plugins
        }

    async def dispatch(
        self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        _starlette_context_token: Token = _request_scope_context_storage.set(
            await self.set_context(request)
        )
        try:
            response = await call_next(request)
            for plugin in self.plugins:
                await plugin.enrich_response(response)

        finally:
            _request_scope_context_storage.reset(_starlette_context_token)

        return response