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)