blob: 587690f036dca4d6b1d1e22a8b2e3366aeb59ac8 (
about) (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
|
# Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
"""
NOTE: All classes and functions in this module are considered private and are
subject to abrupt breaking changes. Please do not use them directly.
"""
from contextlib import contextmanager
from contextvars import ContextVar
from copy import deepcopy
from dataclasses import dataclass, field
from functools import wraps
from typing import Set
@dataclass
class ClientContext:
"""
Encapsulation of objects tracked within the ``_context`` context variable.
``features`` is a set responsible for storing features used during
preparation of an AWS request. ``botocore.useragent.register_feature_id``
is used to add to this set.
"""
features: Set[str] = field(default_factory=set)
_context = ContextVar("_context")
def get_context():
"""Get the current ``_context`` context variable if set, else None."""
return _context.get(None)
def set_context(ctx):
"""Set the current ``_context`` context variable.
:type ctx: ClientContext
:param ctx: Client context object to set as the current context variable.
:rtype: contextvars.Token
:returns: Token object used to revert the context variable to what it was
before the corresponding set.
"""
token = _context.set(ctx)
return token
def reset_context(token):
"""Reset the current ``_context`` context variable.
:type token: contextvars.Token
:param token: Token object to reset the context variable.
"""
_context.reset(token)
@contextmanager
def start_as_current_context(ctx=None):
"""
Context manager that copies the passed or current context object and sets
it as the current context variable. If no context is found, a new
``ClientContext`` object is created. It mainly ensures the context variable
is reset to the previous value once the executed code returns.
Example usage:
def my_feature():
with start_as_current_context():
register_feature_id('MY_FEATURE')
pass
:type ctx: ClientContext
:param ctx: The client context object to set as the new context variable.
If not provided, the current or a new context variable is used.
"""
current = ctx or get_context()
if current is None:
new = ClientContext()
else:
new = deepcopy(current)
token = set_context(new)
try:
yield
finally:
reset_context(token)
def with_current_context(hook=None):
"""
Decorator that wraps ``start_as_current_context`` and optionally invokes a
hook within the newly-set context. This is just syntactic sugar to avoid
indenting existing code under the context manager.
Example usage:
@with_current_context(partial(register_feature_id, 'MY_FEATURE'))
def my_feature():
pass
:type hook: callable
:param hook: A callable that will be invoked within the scope of the
``start_as_current_context`` context manager.
"""
def decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
with start_as_current_context():
if hook:
hook()
return func(*args, **kwargs)
return wrapper
return decorator
|
