8000 chore: support @working_in_progress and @expremental without calling … · ecda909/adk-python@b51a1f4 · GitHub
[go: up one dir, main page]
Skip to content
/ adk-python Public
forked from google/adk-python

Commit b51a1f4

Browse files
seanzhougooglecopybara-github
seanzhougoogle
authored and
copybara-github
committed
chore: support @working_in_progress and @expremental without calling parameters
PiperOrigin-RevId: 771357728
1 parent 8932106 commit b51a1f4

File tree

2 files changed

+184
-62
lines changed

2 files changed

+184
-62
lines changed

‎src/google/adk/utils/feature_decorator.py‎

Lines changed: 93 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -34,70 +34,90 @@ def _make_feature_decorator(
3434
default_message: str,
3535
block_usage: bool = False,
3636
bypass_env_var: Optional[str] = None,
37-
) -> Callable[[str], Callable[[T], T]]:
38-
def decorator_factory(message: str = default_message) -> Callable[[T], T]:
39-
def decorator(obj: T) -> T:
40-
obj_name = getattr(obj, "__name__", type(obj).__name__)
41-
msg = f"[{label.upper()}] {obj_name}: {message}"
42-
43-
if isinstance(obj, type): # decorating a class
44-
orig_init = obj.__init__
45-
46-
@functools.wraps(orig_init)
47-
def new_init(self, *args, **kwargs):
48-
# Load .env file if dotenv is available
49-
load_dotenv()
50-
51-
# Check if usage should be bypassed via environment variable at call time
52-
should_bypass = (
53-
bypass_env_var is not None
54-
and os.environ.get(bypass_env_var, "").lower() == "true"
55-
)
56-
57-
if should_bypass:
58-
# Bypass completely - no warning, no error
59-
pass
60-
elif block_usage:
61-
raise RuntimeError(msg)
62-
else:
63-
warnings.warn(msg, category=UserWarning, stacklevel=2)
64-
return orig_init(self, *args, **kwargs)
65-
66-
obj.__init__ = new_init # type: ignore[attr-defined]
67-
return cast(T, obj)
68-
69-
elif callable(obj): # decorating a function or method
70-
71-
@functools.wraps(obj)
72-
def wrapper(*args, **kwargs):
73-
# Load .env file if dotenv is available
74-
load_dotenv()
75-
76-
# Check if usage should be bypassed via environment variable at call time
77-
should_bypass = (
78-
bypass_env_var is not None
79-
and os.environ.get(bypass_env_var, "").lower() == "true"
80-
)
81-
82-
if should_bypass:
83-
# Bypass completely - no warning, no error
84-
pass
85-
elif block_usage:
86-
raise RuntimeError(msg)
87-
else:
88-
warnings.warn(msg, category=UserWarning, stacklevel=2)
89-
return obj(*args, **kwargs)
90-
91-
return cast(T, wrapper)
92-
93-
else:
94-
raise TypeError(
95-
f"@{label} can only be applied to classes or callable objects"
37+
) -> Callable:
38+
def decorator_factory(message_or_obj=None):
39+
# Case 1: Used as @decorator without parentheses
40+
# message_or_obj is the decorated class/function
41+
if message_or_obj is not None and (
42+
isinstance(message_or_obj, type) or callable(message_or_obj)
43+
):
44+
return _create_decorator(
45+
default_message, label, block_usage, bypass_env_var
46+
)(message_or_obj)
47+
48+
# Case 2: Used as @decorator() with or without message
49+
# message_or_obj is either None or a string message
50+
message = (
51+
message_or_obj if isinstance(message_or_obj, str) else default_message
52+
)
53+
return _create_decorator(message, label, block_usage, bypass_env_var)
54+
55+
return decorator_factory
56+
57+
58+
def _create_decorator(
59+
message: str, label: str, block_usage: bool, bypass_env_var: Optional[str]
60+
) -> Callable[[T], T]:
61+
def decorator(obj: T) -> T:
62+
obj_name = getattr(obj, "__name__", type(obj).__name__)
63+
msg = f"[{label.upper()}] {obj_name}: {message}"
64+
65+
if isinstance(obj, type): # decorating a class
66+
orig_init = obj.__init__
67+
68+
@functools.wraps(orig_init)
69+
def new_init(self, *args, **kwargs):
70+
# Load .env file if dotenv is available
71+
load_dotenv()
72+
73+
# Check if usage should be bypassed via environment variable at call time
74+
should_bypass = (
75+
bypass_env_var is not None
76+
and os.environ.get(bypass_env_var, "").lower() == "true"
9677
)
9778

98-
return decorator
79+
if should_bypass:
80+
# Bypass completely - no warning, no error
81+
pass
82+
elif block_usage:
83+
raise RuntimeError(msg)
84+
else:
85+
warnings.warn(msg, category=UserWarning, stacklevel=2)
86+
return orig_init(self, *args, **kwargs)
87+
88+
obj.__init__ = new_init # type: ignore[attr-defined]
89+
return cast(T, obj)
90+
91+
elif callable(obj): # decorating a function or method
92+
93+
@functools.wraps(obj)
94+
def wrapper(*args, **kwargs):
95+
# Load .env file if dotenv is available
96+
load_dotenv()
97+
98+
# Check if usage should be bypassed via environment variable at call time
99+
should_bypass = (
100+
bypass_env_var is not None
101+
and os.environ.get(bypass_env_var, "").lower() == "true"
102+
)
99103

100-
return decorator_factory
104+
if should_bypass:
105+
# Bypass completely - no warning, no error
106+
pass
107+
elif block_usage:
108+
raise RuntimeError(msg)
109+
else:
110+
warnings.warn(msg, category=UserWarning, stacklevel=2)
111+
return obj(*args, **kwargs)
112+
113+
return cast(T, wrapper)
114+
115+
else:
116+
raise TypeError(
117+
f"@{label} can only be applied to classes or callable objects"
118+
)
119+
120+
return decorator
101121

102122

103123
working_in_progress = _make_feature_decorator(
@@ -137,8 +157,19 @@ def my_wip_function():
137157
Sample usage:
138158
139159
```
140-
@experimental("This API may have breaking change in the future.")
160+
# Use with default message
161+
@experimental
141162
class ExperimentalClass:
142163
pass
164+
165+
# Use with custom message
166+
@experimental("This API may have breaking change in the future.")
167+
class CustomExperimentalClass:
168+
pass
169+
170+
# Use with empty parentheses (same as default message)
171+
@experimental()
172+
def experimental_function():
173+
pass
143174
```
144175
"""

‎tests/unittests/utils/test_feature_decorator.py‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,31 @@ def run(self):
3030
return "running experimental"
3131

3232

33+
# Test classes/functions for new usage patterns
34+
@experimental
35+
class ExperimentalClassNoParens:
36+
37+
def run(self):
38+
return "running experimental without parens"
39+
40+
41+
@experimental()
42+
class ExperimentalClassEmptyParens:
43+
44+
def run(self):
45+
return "running experimental with empty parens"
46+
47+
48+
@experimental
49+
def experimental_fn_no_parens():
50+
return "executing without parens"
51+
52+
53+
@experimental()
54+
def experimental_fn_empty_parens():
55+
return "executing with empty parens"
56+
57+
3358
def test_working_in_progress_class_raises_error():
3459
"""Test that WIP class raises RuntimeError by default."""
3560
# Ensure environment variable is not set
@@ -208,3 +233,69 @@ def test_experimental_class_warns():
208233
assert issubclass(w[0].category, UserWarning)
209234
assert "[EXPERIMENTAL] ExperimentalClass:" in str(w[0].message)
210235
assert "class may change" in str(w[0].message)
236+
237+
238+
def test_experimental_class_no_parens_warns():
239+
"""Test that experimental class without parentheses shows default warning."""
240+
with warnings.catch_warnings(record=True) as w:
241+
warnings.simplefilter("always")
242+
243+
exp_class = ExperimentalClassNoParens()
244+
result = exp_class.run()
245+
246+
assert result == "running experimental without parens"
247+
assert len(w) == 1
248+
assert issubclass(w[0].category, UserWarning)
249+
assert "[EXPERIMENTAL] ExperimentalClassNoParens:" in str(w[0].message)
250+
assert "This feature is experimental and may change or be removed" in str(
251+
w[0].message
252+
)
253+
254+
255+
def test_experimental_class_empty_parens_warns():
256+
"""Test that experimental class with empty parentheses shows default warning."""
257+
with warnings.catch_warnings(record=True) as w:
258+
warnings.simplefilter("always")
259+
260+
exp_class = ExperimentalClassEmptyParens()
261+
result = exp_class.run()
262+
263+
assert result == "running experimental with empty parens"
264+
assert len(w) == 1
265+
assert issubclass(w[0].category, UserWarning)
266+
assert "[EXPERIMENTAL] ExperimentalClassEmptyParens:" in str(w[0].message)
267+
assert "This feature is experimental and may change or be removed" in str(
268+
w[0].message
269+
)
270+
271+
272+
def test_experimental_function_no_parens_warns():
273+
"""Test that experimental function without parentheses shows default warning."""
274+
with warnings.catch_warnings(record=True) as w:
275+
warnings.simplefilter("always")
276+
277+
result = experimental_fn_no_parens()
278+
279+
assert result == "executing without parens"
280+
assert len(w) == 1
281+
assert issubclass(w[0].category, UserWarning)
282+
assert "[EXPERIMENTAL] experimental_fn_no_parens:" in str(w[0].message)
283+
assert "This feature is experimental and may change or be removed" in str(
284+
w[0].message
285+
)
286+
287+
288+
def test_experimental_function_empty_parens_warns():
289+
"""Test that experimental function with empty parentheses shows default warning."""
290+
with warnings.catch_warnings(record=True) as w:
291+
warnings.simplefilter("always")
292+
293+
result = experimental_fn_empty_parens()
294+
295+
assert result == "executing with empty parens"
296+
assert len(w) == 1
297+
assert issubclass(w[0].category, UserWarning)
298+
assert "[EXPERIMENTAL] experimental_fn_empty_parens:" in str(w[0].message)
299+
assert "This feature is experimental and may change or be removed" in str(
300+
w[0].message
301+
)

0 commit comments

Comments
 (0)
0