Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/botocore/context.py: 42%

1# Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"). You 

4# may not use this file except in compliance with the License. A copy of 

5# the License is located at 

6# 

7# http://aws.amazon.com/apache2.0/ 

8# 

9# or in the "license" file accompanying this file. This file is 

10# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF 

11# ANY KIND, either express or implied. See the License for the specific 

12# language governing permissions and limitations under the License. 

13""" 

14NOTE: All classes and functions in this module are considered private and are 

15subject to abrupt breaking changes. Please do not use them directly. 

16""" 

17 

18from contextlib import contextmanager 

19from contextvars import ContextVar 

20from copy import deepcopy 

21from dataclasses import dataclass, field 

22from functools import wraps 

23 

24 

25@dataclass 

26class ClientContext: 

27 """ 

28 Encapsulation of objects tracked within the ``_context`` context variable. 

29 

30 ``features`` is a set responsible for storing features used during 

31 preparation of an AWS request. ``botocore.useragent.register_feature_id`` 

32 is used to add to this set. 

33 """ 

34 

35 features: set[str] = field(default_factory=set) 

36 

37 

38_context = ContextVar("_context") 

39 

40 

41def get_context(): 

42 """Get the current ``_context`` context variable if set, else None.""" 

43 return _context.get(None) 

44 

45 

46def set_context(ctx): 

47 """Set the current ``_context`` context variable. 

48 

49 :type ctx: ClientContext 

50 :param ctx: Client context object to set as the current context variable. 

51 

52 :rtype: contextvars.Token 

53 :returns: Token object used to revert the context variable to what it was 

54 before the corresponding set. 

55 """ 

56 token = _context.set(ctx) 

57 return token 

58 

59 

60def reset_context(token): 

61 """Reset the current ``_context`` context variable. 

62 

63 :type token: contextvars.Token 

64 :param token: Token object to reset the context variable. 

65 """ 

66 _context.reset(token) 

67 

68 

69@contextmanager 

70def start_as_current_context(ctx=None): 

71 """ 

72 Context manager that copies the passed or current context object and sets 

73 it as the current context variable. If no context is found, a new 

74 ``ClientContext`` object is created. It mainly ensures the context variable 

75 is reset to the previous value once the executed code returns. 

76 

77 Example usage: 

78 

79 def my_feature(): 

80 with start_as_current_context(): 

81 register_feature_id('MY_FEATURE') 

82 pass 

83 

84 :type ctx: ClientContext 

85 :param ctx: The client context object to set as the new context variable. 

86 If not provided, the current or a new context variable is used. 

87 """ 

88 current = ctx or get_context() 

89 if current is None: 

90 new = ClientContext() 

91 else: 

92 new = deepcopy(current) 

93 token = set_context(new) 

94 try: 

95 yield 

96 finally: 

97 reset_context(token) 

98 

99 

100def with_current_context(hook=None): 

101 """ 

102 Decorator that wraps ``start_as_current_context`` and optionally invokes a 

103 hook within the newly-set context. This is just syntactic sugar to avoid 

104 indenting existing code under the context manager. 

105 

106 Example usage: 

107 

108 @with_current_context(partial(register_feature_id, 'MY_FEATURE')) 

109 def my_feature(): 

110 pass 

111 

112 :type hook: callable 

113 :param hook: A callable that will be invoked within the scope of the 

114 ``start_as_current_context`` context manager. 

115 """ 

116 

117 def decorator(func): 

118 @wraps(func) 

119 def wrapper(*args, **kwargs): 

120 with start_as_current_context(): 

121 if hook: 

122 hook() 

123 return func(*args, **kwargs) 

124 

125 return wrapper 

126 

127 return decorator