Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/airflow/utils/deprecation_tools.py: 36%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

50 statements  

1# Licensed to the Apache Software Foundation (ASF) under one 

2# or more contributor license agreements. See the NOTICE file 

3# distributed with this work for additional information 

4# regarding copyright ownership. The ASF licenses this file 

5# to you under the Apache License, Version 2.0 (the 

6# "License"); you may not use this file except in compliance 

7# with the License. You may obtain a copy of the License at 

8# 

9# http://www.apache.org/licenses/LICENSE-2.0 

10# 

11# Unless required by applicable law or agreed to in writing, 

12# software distributed under the License is distributed on an 

13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 

14# KIND, either express or implied. See the License for the 

15# specific language governing permissions and limitations 

16# under the License. 

17from __future__ import annotations 

18 

19import functools 

20import importlib 

21import sys 

22import warnings 

23from types import ModuleType 

24 

25 

26class DeprecatedImportWarning(FutureWarning): 

27 """ 

28 Warning class for deprecated imports in Airflow. 

29 

30 This warning is raised when users import deprecated classes or functions 

31 from Airflow modules that have been moved to better locations. 

32 """ 

33 

34 ... 

35 

36 

37def getattr_with_deprecation( 

38 imports: dict[str, str], 

39 module: str, 

40 override_deprecated_classes: dict[str, str], 

41 extra_message: str, 

42 name: str, 

43 message_override: str = "", 

44): 

45 """ 

46 Retrieve the imported attribute from the redirected module and raises a deprecation warning. 

47 

48 :param imports: dict of imports and their redirection for the module 

49 :param module: name of the module in the package to get the attribute from 

50 :param override_deprecated_classes: override target attributes with deprecated ones. If target attribute is 

51 found in the dictionary, it will be displayed in the warning message. 

52 :param extra_message: extra message to display in the warning or import error message 

53 :param message_override: if provided, overrides the default deprecation message. Supports placeholders: 

54 {module}, {name}, {target} which are substituted with the actual values. 

55 :param name: attribute name 

56 :return: 

57 """ 

58 target_class_full_name = imports.get(name) 

59 

60 # Handle wildcard pattern "*" - redirect all attributes to target module 

61 # Skip Python special attributes (dunder attributes) as they shouldn't be redirected 

62 if not target_class_full_name and "*" in imports and not (name.startswith("__") and name.endswith("__")): 

63 target_class_full_name = f"{imports['*']}.{name}" 

64 

65 if not target_class_full_name: 

66 raise AttributeError(f"The module `{module!r}` has no attribute `{name!r}`") 

67 

68 # Determine the warning class name (may be overridden) 

69 warning_class_name = target_class_full_name 

70 if override_deprecated_classes and name in override_deprecated_classes: 

71 warning_class_name = override_deprecated_classes[name] 

72 

73 if message_override: 

74 message = message_override.format(module=module, name=name, target=warning_class_name) 

75 else: 

76 message = f"The `{module}.{name}` attribute is deprecated. Please use `{warning_class_name!r}`." 

77 if extra_message: 

78 message += f" {extra_message}." 

79 warnings.warn(message, DeprecatedImportWarning, stacklevel=2) 

80 

81 # Import and return the target attribute 

82 new_module, new_class_name = target_class_full_name.rsplit(".", 1) 

83 try: 

84 return getattr(importlib.import_module(new_module), new_class_name) 

85 except ImportError as e: 

86 error_message = ( 

87 f"Could not import `{new_module}.{new_class_name}` while trying to import `{module}.{name}`." 

88 ) 

89 if extra_message: 

90 error_message += f" {extra_message}." 

91 raise ImportError(error_message) from e 

92 

93 

94def add_deprecated_classes( 

95 module_imports: dict[str, dict[str, str]], 

96 package: str, 

97 override_deprecated_classes: dict[str, dict[str, str]] | None = None, 

98 extra_message: str | None = None, 

99 message: str | None = None, 

100): 

101 """ 

102 Add deprecated attribute PEP-563 imports and warnings modules to the package. 

103 

104 Works for classes, functions, variables, and other module attributes. 

105 Supports both creating virtual modules and modifying existing modules. 

106 

107 :param module_imports: imports to use. Format: dict[str, dict[str, str]] 

108 - Keys are module names (creates virtual modules) 

109 - Special key __name__ modifies the current module for direct attribute imports 

110 - Can mix both approaches in a single call 

111 :param package: package name (typically __name__) 

112 :param override_deprecated_classes: override target attributes with deprecated ones. 

113 Format: dict[str, dict[str, str]] matching the structure of module_imports 

114 :param extra_message: extra message to display in the warning or import error message 

115 :param message: if provided, overrides the default deprecation message. Supports placeholders: 

116 {module}, {name}, {target} which are substituted with the actual values. 

117 

118 Examples: 

119 # Create virtual modules (e.g., for removed .py files) 

120 add_deprecated_classes( 

121 {"basenotifier": {"BaseNotifier": "airflow.sdk.bases.notifier.BaseNotifier"}}, 

122 package=__name__, 

123 ) 

124 

125 # Wildcard support - redirect all attributes to new module 

126 add_deprecated_classes( 

127 {"timezone": {"*": "airflow.sdk.timezone"}}, 

128 package=__name__, 

129 ) 

130 

131 # Current module direct imports 

132 add_deprecated_classes( 

133 { 

134 __name__: { 

135 "get_fs": "airflow.sdk.io.fs.get_fs", 

136 "has_fs": "airflow.sdk.io.fs.has_fs", 

137 } 

138 }, 

139 package=__name__, 

140 ) 

141 

142 # Mixed behavior - both current module and submodule attributes 

143 add_deprecated_classes( 

144 { 

145 __name__: { 

146 "get_fs": "airflow.sdk.io.fs.get_fs", 

147 "has_fs": "airflow.sdk.io.fs.has_fs", 

148 "Properties": "airflow.sdk.io.typedef.Properties", 

149 }, 

150 "typedef": { 

151 "Properties": "airflow.sdk.io.typedef.Properties", 

152 } 

153 }, 

154 package=__name__, 

155 ) 

156 

157 The first example makes 'from airflow.notifications.basenotifier import BaseNotifier' work 

158 even if 'basenotifier.py' was removed. 

159 

160 The second example makes 'from airflow.utils.timezone import utc' redirect to 'airflow.sdk.timezone.utc', 

161 allowing any attribute from the deprecated module to be accessed from the new location. 

162 

163 The third example makes 'from airflow.io import get_fs' work with direct imports from the current module. 

164 

165 The fourth example handles both direct imports from the current module and submodule imports. 

166 """ 

167 # Handle both current module and virtual module deprecations 

168 for module_name, imports in module_imports.items(): 

169 if module_name == package: 

170 # Special case: modify the current module for direct attribute imports 

171 if package not in sys.modules: 

172 raise ValueError(f"Module {package} not found in sys.modules") 

173 

174 module = sys.modules[package] 

175 

176 # Create the __getattr__ function for current module 

177 current_override = {} 

178 if override_deprecated_classes and package in override_deprecated_classes: 

179 current_override = override_deprecated_classes[package] 

180 

181 getattr_func = functools.partial( 

182 getattr_with_deprecation, 

183 imports, 

184 package, 

185 current_override, 

186 extra_message or "", 

187 message_override=message or "", 

188 ) 

189 

190 # Set the __getattr__ function on the current module 

191 setattr(module, "__getattr__", getattr_func) 

192 else: 

193 # Create virtual modules for submodule imports 

194 full_module_name = f"{package}.{module_name}" 

195 module_type = ModuleType(full_module_name) 

196 if override_deprecated_classes and module_name in override_deprecated_classes: 

197 override_deprecated_classes_for_module = override_deprecated_classes[module_name] 

198 else: 

199 override_deprecated_classes_for_module = {} 

200 

201 # Mypy is not able to derive the right function signature https://github.com/python/mypy/issues/2427 

202 module_type.__getattr__ = functools.partial( # type: ignore[method-assign] 

203 getattr_with_deprecation, 

204 imports, 

205 full_module_name, 

206 override_deprecated_classes_for_module, 

207 extra_message or "", 

208 message_override=message or "", 

209 ) 

210 sys.modules.setdefault(full_module_name, module_type)