Source code for hbutils.reflection.context

"""
Overview:
    Utilities for building context variables on thread level.

    This is useful when implementing a with-block-based syntax.
    For example:

        >>> from contextlib import contextmanager
        >>> from hbutils.reflection import context
        >>>
        >>> # developer's view
        ... @contextmanager
        ... def use_mul():  # set 'mul' to `True` in its with-block
        ...     with context().vars(mul=True):
        ...         yield
        >>>
        >>> def calc(a, b):  # logic of `calc` will be changed when 'mul' is given
        ...     if context().get('mul', None):
        ...         return a * b
        ...     else:
        ...         return a + b
        >>>
        >>> # user's view (magic-liked, isn't it?)
        ... print(calc(3, 5))  # 3 + 5
        8
        >>> with use_mul():
        ...     print(calc(3, 5))  # changed to 3 * 5
        15
        >>> print(calc(3, 5))  # back to 3 + 5, again :)
        8
"""
import collections.abc

from contextlib import contextmanager
from functools import wraps
from multiprocessing import current_process
from threading import current_thread
from typing import Tuple, TypeVar, Iterator, Mapping, Optional

__all__ = [
    'context', 'cwrap',
]


def _get_pid() -> int:
    return current_process().pid


def _get_tid() -> int:
    return current_thread().ident


def _get_context_id() -> Tuple[int, int]:
    return _get_pid(), _get_tid()


_global_contexts = {}

_KeyType = TypeVar('_KeyType', bound=str)
_ValueType = TypeVar('_ValueType')


[docs]class ContextVars(collections.abc.Mapping): """ Overview: Context variable management. .. note:: This class is inherited from :class:`collections.abc.Mapping`. Main features of mapping object (such as ``__getitem__``, ``__len__``, ``__iter__``) are supported. See `Collections Abstract Base Classes \ <https://docs.python.org/3/library/collections.abc.html#collections-abstract-base-classes>`_. .. warning:: This object should be singleton on thread level. It is not recommended constructing manually. """
[docs] def __init__(self, **kwargs): """ Constructor of :class:`ContextVars`. :param kwargs: Initial context values. """ self._vars = dict(kwargs)
@contextmanager def _with_vars(self, params: Mapping[_KeyType, _ValueType], clear: bool = False): # initialize new values _origin = dict(self._vars) self._vars.update(params) if clear: for key in list(self._vars.keys()): if key not in params: del self._vars[key] try: yield finally: # de-initialize, recover changed values for k in set(_origin.keys()) | set(self._vars.keys()): if k not in _origin: del self._vars[k] else: self._vars[k] = _origin[k]
[docs] @contextmanager def vars(self, **kwargs): """ Adding variables into context of ``with`` block. :param kwargs: Additional context variables. Examples:: >>> from hbutils.reflection import context >>> >>> def var_detect(): ... if context().get('var', None): ... print(f'Var detected, its value is {context()["var"]}.') ... else: ... print('Var not detected.') >>> >>> var_detect() Var not detected. >>> with context().vars(var=1): ... var_detect() Var detected, its value is 1. >>> var_detect() Var not detected. .. note:: See :func:`context`. """ with self._with_vars(kwargs, clear=False): yield
[docs] @contextmanager def inherit(self, context_: 'ContextVars'): """ Inherit variables from the given context. :param context_: :class:`ContextVars` object to inherit from. .. note:: After :meth:`inherit` is used, **the original variables which not present in the given ``context_`` \ will be removed**. This is different from :meth:`vars`, so attention. """ with self._with_vars(context_._vars, clear=True): yield
def __getitem__(self, key: _KeyType): return self._vars[key] def __len__(self) -> int: return len(self._vars) def __iter__(self) -> Iterator[_KeyType]: return self._vars.__iter__()
[docs]def context() -> ContextVars: """ Overview: Get context object in this thread. :return: Context object in this thread. .. note:: This result is unique on one thread. """ _context_id = _get_context_id() if _context_id not in _global_contexts: _context = ContextVars() _global_contexts[_context_id] = _context return _global_contexts[_context_id]
[docs]def cwrap(func, *, context_: Optional[ContextVars] = None, **vars_): """ Overview: Context wrap for functions. :param func: Original function to wrap. :param context_: Context for inheriting. Default is ``None`` which means :func:`context`'s result will be used. :param vars_: External variables after inherit context. .. note:: :func:`cwrap` is important when you need to pass the current context into thread. And **it is compitable on all platforms**. For example: >>> from threading import Thread >>> from hbutils.reflection import context, cwrap >>> >>> def var_detect(): ... if context().get('var', None): ... print(f'Var detected, its value is {context()["var"]}.') ... else: ... print('Var not detected.') >>> >>> with context().vars(var=1): # no inherit, vars will be lost in thread ... t = Thread(target=var_detect) ... t.start() ... t.join() Var not detected. >>> with context().vars(var=1): # with inherit, vars will be kept in thread ... t = Thread(target=cwrap(var_detect)) ... t.start() ... t.join() Var detected, its value is 1. .. warning:: :func:`cwrap` **is not compitable on Windows or Python3.8+ on macOS** when creating **new process**. Please pass in direct arguments by ``args`` argument of :class:`Process`. If you insist on using :func:`context` feature, you need to pass the context object into the sub process. For example: >>> from contextlib import contextmanager >>> from multiprocessing import Process >>> from hbutils.reflection import context >>> >>> # developer's view ... @contextmanager ... def use_mul(): # set 'mul' to `True` in its with-block ... with context().vars(mul=True): ... yield >>> >>> def calc(a, b): # logic of `calc` will be changed when 'mul' is given ... if context().get('mul', None): ... print(a * b) ... else: ... print(a + b) >>> >>> def _calc(a, b, ctx=None): ... with context().inherit(ctx or context()): ... return calc(a, b) >>> >>> # user's view ... if __name__ == '__main__': ... calc(3, 5) # 3 + 5 ... with use_mul(): ... p = Process(target=_calc, args=(3, 5, context())) ... p.start() ... p.join() ... calc(3, 5) # back to 3 + 5, again :) 8 15 8 """ context_ = context_ or context() @wraps(func) def _new_func(*args, **kwargs): with context().inherit(context_): with context().vars(**vars_): return func(*args, **kwargs) return _new_func