Skip to content

Commit 4b6c42b

Browse files
authored
Extend add_deprecated_classes to support wildcard patterns (#53196)
Add support for wildcard pattern `"*"` in `add_deprecated_classes` function to redirect any attribute access from a deprecated module to a target module. Changes: - Extended `getattr_with_deprecation` to handle `"*"` wildcard pattern - Added comprehensive test suite with side-effect-free test patterns - Updated documentation with wildcard usage examples - Added testing standards for avoiding side effects between tests This enables patterns like: ```py add_deprecated_classes({ "timezone": {"": "airflow.sdk.timezone"} }, package=name) ``` Where any import from the deprecated module gets redirected to the new location with appropriate deprecation warnings. The implementation maintains backward compatibility with existing specific class mappings taking priority over wildcard redirects.
1 parent bc79040 commit 4b6c42b

File tree

2 files changed

+348
-6
lines changed

2 files changed

+348
-6
lines changed

airflow-core/src/airflow/utils/deprecation_tools.py

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -35,22 +35,33 @@ def getattr_with_deprecation(
3535
3636
:param imports: dict of imports and their redirection for the module
3737
:param module: name of the module in the package to get the attribute from
38-
:param override_deprecated_classes: override target classes with deprecated ones. If target class is
38+
:param override_deprecated_classes: override target attributes with deprecated ones. If target attribute is
3939
found in the dictionary, it will be displayed in the warning message.
4040
:param extra_message: extra message to display in the warning or import error message
4141
:param name: attribute name
4242
:return:
4343
"""
4444
target_class_full_name = imports.get(name)
45+
46+
# Handle wildcard pattern "*" - redirect all attributes to target module
47+
# Skip Python special attributes (dunder attributes) as they shouldn't be redirected
48+
if not target_class_full_name and "*" in imports and not (name.startswith("__") and name.endswith("__")):
49+
target_class_full_name = f"{imports['*']}.{name}"
50+
4551
if not target_class_full_name:
4652
raise AttributeError(f"The module `{module!r}` has no attribute `{name!r}`")
53+
54+
# Determine the warning class name (may be overridden)
4755
warning_class_name = target_class_full_name
4856
if override_deprecated_classes and name in override_deprecated_classes:
4957
warning_class_name = override_deprecated_classes[name]
50-
message = f"The `{module}.{name}` class is deprecated. Please use `{warning_class_name!r}`."
58+
59+
message = f"The `{module}.{name}` attribute is deprecated. Please use `{warning_class_name!r}`."
5160
if extra_message:
5261
message += f" {extra_message}."
5362
warnings.warn(message, DeprecationWarning, stacklevel=2)
63+
64+
# Import and return the target attribute
5465
new_module, new_class_name = target_class_full_name.rsplit(".", 1)
5566
try:
5667
return getattr(importlib.import_module(new_module), new_class_name)
@@ -70,14 +81,14 @@ def add_deprecated_classes(
7081
extra_message: str | None = None,
7182
):
7283
"""
73-
Add deprecated class PEP-563 imports and warnings modules to the package.
84+
Add deprecated attribute PEP-563 imports and warnings modules to the package.
7485
75-
Side note: It also works for methods, not just classes.
86+
Works for classes, functions, variables, and other module attributes.
7687
7788
:param module_imports: imports to use
7889
:param package: package name
79-
:param override_deprecated_classes: override target classes with deprecated ones. If module +
80-
target class is found in the dictionary, it will be displayed in the warning message.
90+
:param override_deprecated_classes: override target attributes with deprecated ones. If module +
91+
target attribute is found in the dictionary, it will be displayed in the warning message.
8192
:param extra_message: extra message to display in the warning or import error message
8293
8394
Example:
@@ -89,6 +100,15 @@ def add_deprecated_classes(
89100
This makes 'from airflow.notifications.basenotifier import BaseNotifier' still work,
90101
even if 'basenotifier.py' was removed, and shows a warning with the new path.
91102
103+
Wildcard Example:
104+
add_deprecated_classes(
105+
{"timezone": {"*": "airflow.sdk.timezone"}},
106+
package=__name__,
107+
)
108+
109+
This makes 'from airflow.utils.timezone import utc' redirect to 'airflow.sdk.timezone.utc',
110+
allowing any attribute from the deprecated module to be accessed from the new location.
111+
92112
Note that "add_deprecated_classes method should be called in the `__init__.py` file in the package
93113
where the deprecated classes are located - this way the module `.py` files should be removed and what
94114
remains in the package is just the `__init__.py` file.
Lines changed: 322 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,322 @@
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.
17+
from __future__ import annotations
18+
19+
import contextlib
20+
import re
21+
import sys
22+
import uuid
23+
import warnings
24+
from contextlib import contextmanager
25+
from types import ModuleType
26+
from unittest import mock
27+
28+
import pytest
29+
30+
from airflow.utils.deprecation_tools import add_deprecated_classes, getattr_with_deprecation
31+
32+
33+
@contextmanager
34+
def temporary_module(module_name):
35+
"""Context manager to safely add and remove modules from sys.modules."""
36+
original_module = sys.modules.get(module_name)
37+
try:
38+
yield
39+
finally:
40+
if original_module is not None:
41+
sys.modules[module_name] = original_module
42+
elif module_name in sys.modules:
43+
del sys.modules[module_name]
44+
45+
46+
def get_unique_module_name(base_name="test_module"):
47+
"""Generate a unique module name to avoid conflicts."""
48+
return f"{base_name}_{uuid.uuid4().hex[:8]}"
49+
50+
51+
class TestGetAttrWithDeprecation:
52+
"""Tests for the getattr_with_deprecation function."""
53+
54+
def test_getattr_with_deprecation_specific_class(self):
55+
"""Test deprecated import for a specific class."""
56+
imports = {"OldClass": "new.module.NewClass"}
57+
58+
# Mock the new module and class
59+
mock_module = mock.MagicMock()
60+
mock_new_class = mock.MagicMock()
61+
mock_module.NewClass = mock_new_class
62+
63+
with mock.patch("airflow.utils.deprecation_tools.importlib.import_module", return_value=mock_module):
64+
with warnings.catch_warnings(record=True) as w:
65+
warnings.simplefilter("always")
66+
result = getattr_with_deprecation(
67+
imports=imports,
68+
module="old.module",
69+
override_deprecated_classes={},
70+
extra_message="",
71+
name="OldClass",
72+
)
73+
74+
assert result == mock_new_class
75+
assert len(w) == 1
76+
assert issubclass(w[0].category, DeprecationWarning)
77+
assert "old.module.OldClass" in str(w[0].message)
78+
assert "new.module.NewClass" in str(w[0].message)
79+
80+
def test_getattr_with_deprecation_wildcard(self):
81+
"""Test deprecated import using wildcard pattern."""
82+
imports = {"*": "new.module"}
83+
84+
# Mock the new module and attribute
85+
mock_module = mock.MagicMock()
86+
mock_attribute = mock.MagicMock()
87+
mock_module.SomeAttribute = mock_attribute
88+
89+
with mock.patch("airflow.utils.deprecation_tools.importlib.import_module", return_value=mock_module):
90+
with warnings.catch_warnings(record=True) as w:
91+
warnings.simplefilter("always")
92+
result = getattr_with_deprecation(
93+
imports=imports,
94+
module="old.module",
95+
override_deprecated_classes={},
96+
extra_message="",
97+
name="SomeAttribute",
98+
)
99+
100+
assert result == mock_attribute
101+
assert len(w) == 1
102+
assert issubclass(w[0].category, DeprecationWarning)
103+
assert "old.module.SomeAttribute" in str(w[0].message)
104+
assert "new.module.SomeAttribute" in str(w[0].message)
105+
106+
def test_getattr_with_deprecation_wildcard_with_override(self):
107+
"""Test wildcard pattern with override deprecated classes."""
108+
imports = {"*": "new.module"}
109+
override_deprecated_classes = {"SomeAttribute": "override.module.OverrideClass"}
110+
111+
# Mock the new module and attribute
112+
mock_module = mock.MagicMock()
113+
mock_attribute = mock.MagicMock()
114+
mock_module.SomeAttribute = mock_attribute
115+
116+
with mock.patch("airflow.utils.deprecation_tools.importlib.import_module", return_value=mock_module):
117+
with warnings.catch_warnings(record=True) as w:
118+
warnings.simplefilter("always")
119+
result = getattr_with_deprecation(
120+
imports=imports,
121+
module="old.module",
122+
override_deprecated_classes=override_deprecated_classes,
123+
extra_message="",
124+
name="SomeAttribute",
125+
)
126+
127+
assert result == mock_attribute
128+
assert len(w) == 1
129+
assert issubclass(w[0].category, DeprecationWarning)
130+
assert "old.module.SomeAttribute" in str(w[0].message)
131+
assert "override.module.OverrideClass" in str(w[0].message)
132+
133+
def test_getattr_with_deprecation_specific_class_priority(self):
134+
"""Test that specific class mapping takes priority over wildcard."""
135+
imports = {"SpecificClass": "specific.module.SpecificClass", "*": "wildcard.module"}
136+
137+
# Mock the specific module and class
138+
mock_module = mock.MagicMock()
139+
mock_specific_class = mock.MagicMock()
140+
mock_module.SpecificClass = mock_specific_class
141+
142+
with mock.patch("airflow.utils.deprecation_tools.importlib.import_module", return_value=mock_module):
143+
with warnings.catch_warnings(record=True) as w:
144+
warnings.simplefilter("always")
145+
result = getattr_with_deprecation(
146+
imports=imports,
147+
module="old.module",
148+
override_deprecated_classes={},
149+
extra_message="",
150+
name="SpecificClass",
151+
)
152+
153+
assert result == mock_specific_class
154+
assert len(w) == 1
155+
assert issubclass(w[0].category, DeprecationWarning)
156+
assert "old.module.SpecificClass" in str(w[0].message)
157+
assert "specific.module.SpecificClass" in str(w[0].message)
158+
159+
def test_getattr_with_deprecation_attribute_not_found(self):
160+
"""Test AttributeError when attribute not found."""
161+
imports = {"ExistingClass": "new.module.ExistingClass"}
162+
163+
with pytest.raises(AttributeError, match=r"has no attribute.*NonExistentClass"):
164+
getattr_with_deprecation(
165+
imports=imports,
166+
module="old.module",
167+
override_deprecated_classes={},
168+
extra_message="",
169+
name="NonExistentClass",
170+
)
171+
172+
def test_getattr_with_deprecation_import_error(self):
173+
"""Test ImportError when target module cannot be imported."""
174+
imports = {"*": "nonexistent.module"}
175+
176+
with mock.patch(
177+
"airflow.utils.deprecation_tools.importlib.import_module",
178+
side_effect=ImportError("Module not found"),
179+
):
180+
with pytest.raises(ImportError, match="Could not import"):
181+
getattr_with_deprecation(
182+
imports=imports,
183+
module="old.module",
184+
override_deprecated_classes={},
185+
extra_message="",
186+
name="SomeAttribute",
187+
)
188+
189+
def test_getattr_with_deprecation_with_extra_message(self):
190+
"""Test that extra message is included in warning."""
191+
imports = {"*": "new.module"}
192+
extra_message = "This is an extra message"
193+
194+
# Mock the new module and attribute
195+
mock_module = mock.MagicMock()
196+
mock_attribute = mock.MagicMock()
197+
mock_module.SomeAttribute = mock_attribute
198+
199+
with mock.patch("airflow.utils.deprecation_tools.importlib.import_module", return_value=mock_module):
200+
with warnings.catch_warnings(record=True) as w:
201+
warnings.simplefilter("always")
202+
getattr_with_deprecation(
203+
imports=imports,
204+
module="old.module",
205+
override_deprecated_classes={},
206+
extra_message=extra_message,
207+
name="SomeAttribute",
208+
)
209+
210+
assert len(w) == 1
211+
assert extra_message in str(w[0].message)
212+
213+
@pytest.mark.parametrize("dunder_attribute", ["__path__", "__file__"])
214+
def test_getattr_with_deprecation_wildcard_skips_dunder_attributes(self, dunder_attribute):
215+
"""Test that wildcard pattern skips Python special attributes."""
216+
imports = {"*": "new.module"}
217+
218+
# Special attributes should raise AttributeError, not be redirected
219+
with pytest.raises(AttributeError, match=rf"has no attribute.*{re.escape(dunder_attribute)}"):
220+
getattr_with_deprecation(
221+
imports=imports,
222+
module="old.module",
223+
override_deprecated_classes={},
224+
extra_message="",
225+
name=dunder_attribute,
226+
)
227+
228+
@pytest.mark.parametrize("non_dunder_attr", ["__version", "__author", "_private", "public"])
229+
def test_getattr_with_deprecation_wildcard_allows_non_dunder_attributes(self, non_dunder_attr):
230+
"""Test that wildcard pattern allows non-dunder attributes (including single underscore prefixed)."""
231+
imports = {"*": "unittest.mock"}
232+
233+
# These should be redirected through wildcard pattern
234+
with warnings.catch_warnings(record=True) as w:
235+
warnings.simplefilter("always")
236+
with contextlib.suppress(ImportError, AttributeError):
237+
# Expected - the target module might not have the attribute
238+
# The important thing is that it tried to redirect (didn't raise AttributeError immediately)
239+
getattr_with_deprecation(
240+
imports=imports,
241+
module="old.module",
242+
override_deprecated_classes={},
243+
extra_message="",
244+
name=non_dunder_attr,
245+
)
246+
247+
# Should have generated a deprecation warning
248+
assert len(w) == 1
249+
assert "deprecated" in str(w[0].message).lower()
250+
251+
252+
class TestAddDeprecatedClasses:
253+
"""Tests for the add_deprecated_classes function."""
254+
255+
def test_add_deprecated_classes_basic(self):
256+
"""Test basic functionality of add_deprecated_classes."""
257+
# Use unique package and module names to avoid conflicts
258+
package_name = get_unique_module_name("test_package")
259+
module_name = f"{package_name}.old_module"
260+
261+
module_imports = {"old_module": {"OldClass": "new.module.NewClass"}}
262+
263+
with temporary_module(module_name):
264+
add_deprecated_classes(module_imports, package_name)
265+
266+
# Check that the module was added to sys.modules
267+
assert module_name in sys.modules
268+
assert isinstance(sys.modules[module_name], ModuleType)
269+
assert hasattr(sys.modules[module_name], "__getattr__")
270+
271+
def test_add_deprecated_classes_with_wildcard(self):
272+
"""Test add_deprecated_classes with wildcard pattern."""
273+
# Use unique package and module names to avoid conflicts
274+
package_name = get_unique_module_name("test_package")
275+
module_name = f"{package_name}.timezone"
276+
277+
module_imports = {"timezone": {"*": "airflow.sdk.timezone"}}
278+
279+
with temporary_module(module_name):
280+
add_deprecated_classes(module_imports, package_name)
281+
282+
# Check that the module was added to sys.modules
283+
assert module_name in sys.modules
284+
assert isinstance(sys.modules[module_name], ModuleType)
285+
assert hasattr(sys.modules[module_name], "__getattr__")
286+
287+
def test_add_deprecated_classes_with_override(self):
288+
"""Test add_deprecated_classes with override_deprecated_classes."""
289+
# Use unique package and module names to avoid conflicts
290+
package_name = get_unique_module_name("test_package")
291+
module_name = f"{package_name}.old_module"
292+
293+
module_imports = {"old_module": {"OldClass": "new.module.NewClass"}}
294+
295+
override_deprecated_classes = {"old_module": {"OldClass": "override.module.OverrideClass"}}
296+
297+
with temporary_module(module_name):
298+
add_deprecated_classes(module_imports, package_name, override_deprecated_classes)
299+
300+
# Check that the module was added to sys.modules
301+
assert module_name in sys.modules
302+
assert isinstance(sys.modules[module_name], ModuleType)
303+
304+
def test_add_deprecated_classes_doesnt_override_existing(self):
305+
"""Test that add_deprecated_classes doesn't override existing modules."""
306+
# Use unique package and module names to avoid conflicts
307+
package_name = get_unique_module_name("test_package")
308+
module_name = f"{package_name}.existing_module"
309+
310+
module_imports = {"existing_module": {"SomeClass": "new.module.SomeClass"}}
311+
312+
with temporary_module(module_name):
313+
# Create a mock existing module
314+
existing_module = ModuleType(module_name)
315+
existing_module.existing_attribute = "existing_value"
316+
sys.modules[module_name] = existing_module
317+
318+
add_deprecated_classes(module_imports, package_name)
319+
320+
# Check that the existing module was not overridden
321+
assert sys.modules[module_name] is existing_module
322+
assert sys.modules[module_name].existing_attribute == "existing_value"

0 commit comments

Comments
 (0)