Skip to content

Commit 4afc9b2

Browse files
George Wealecopybara-github
authored andcommitted
feat: Add env var to suppress experimental warnings
Checked with local wheel and worked as intended. The harness shows suppression works: 0 warnings for all true-like values. This CL adds ADK_DISABLE_EXPERIMENTAL_WARNING to let the users to suppress warning messages from features decorated with @experimental. Previously, using experimental features would always trigger a UserWarning. This change creates a way to disable these warnings, which can be good to stop flooding logs. The warning is suppressed if ADK_DISABLE_EXPERIMENTAL_WARNING is set to a truthy value such as "true", "1", "yes", or "on" (case-insensitive). Added unit tests to make sure: Warning suppression for functions and classes when the env var is set. Case-insensitivity and various truthy values for the env var. Loading the env var from a .env file. PiperOrigin-RevId: 796649404
1 parent f8fd6a4 commit 4afc9b2

File tree

2 files changed

+71
-8
lines changed

2 files changed

+71
-8
lines changed

src/google/adk/utils/feature_decorator.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,13 @@
2626
T = TypeVar("T", bound=Union[Callable, type])
2727

2828

29+
def _is_truthy_env(var_name: str) -> bool:
30+
value = os.environ.get(var_name)
31+
if value is None:
32+
return False
33+
return value.strip().lower() in ("1", "true", "yes", "on")
34+
35+
2936
def _make_feature_decorator(
3037
*,
3138
label: str,
@@ -66,9 +73,8 @@ def decorator(obj: T) -> T:
6673
@functools.wraps(orig_init)
6774
def new_init(self, *args, **kwargs):
6875
# Check if usage should be bypassed via environment variable at call time
69-
should_bypass = (
70-
bypass_env_var is not None
71-
and os.environ.get(bypass_env_var, "").lower() == "true"
76+
should_bypass = bypass_env_var is not None and _is_truthy_env(
77+
bypass_env_var
7278
)
7379

7480
if should_bypass:
@@ -88,9 +94,8 @@ def new_init(self, *args, **kwargs):
8894
@functools.wraps(obj)
8995
def wrapper(*args, **kwargs):
9096
# Check if usage should be bypassed via environment variable at call time
91-
should_bypass = (
92-
bypass_env_var is not None
93-
and os.environ.get(bypass_env_var, "").lower() == "true"
97+
should_bypass = bypass_env_var is not None and _is_truthy_env(
98+
bypass_env_var
9499
)
95100

96101
if should_bypass:
@@ -143,6 +148,7 @@ def my_wip_function():
143148
" versions without notice. It may introduce breaking changes at any"
144149
" time."
145150
),
151+
bypass_env_var="ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS",
146152
)
147153
"""Mark a class or a function as an experimental feature.
148154

tests/unittests/utils/test_feature_decorator.py

Lines changed: 59 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -206,8 +206,12 @@ def test_working_in_progress_loads_from_dotenv_file():
206206
del os.environ["ADK_ALLOW_WIP_FEATURES"]
207207

208208

209-
def test_experimental_function_warns():
209+
def test_experimental_function_warns(monkeypatch):
210210
"""Test that experimental function shows warnings (unchanged behavior)."""
211+
# Ensure environment variable is not set
212+
monkeypatch.delenv(
213+
"ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", raising=False
214+
)
211215
with warnings.catch_warnings(record=True) as w:
212216
warnings.simplefilter("always")
213217

@@ -220,8 +224,12 @@ def test_experimental_function_warns():
220224
assert "breaking change in the future" in str(w[0].message)
221225

222226

223-
def test_experimental_class_warns():
227+
def test_experimental_class_warns(monkeypatch):
224228
"""Test that experimental class shows warnings (unchanged behavior)."""
229+
# Ensure environment variable is not set
230+
monkeypatch.delenv(
231+
"ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", raising=False
232+
)
225233
with warnings.catch_warnings(record=True) as w:
226234
warnings.simplefilter("always")
227235

@@ -235,6 +243,55 @@ def test_experimental_class_warns():
235243
assert "class may change" in str(w[0].message)
236244

237245

246+
def test_experimental_function_bypassed_with_env_var(monkeypatch):
247+
"""Experimental function emits no warning when bypass env var is true."""
248+
true_values = ["true", "True", "TRUE", "1", "yes", "YES", "on", "ON"]
249+
for true_val in true_values:
250+
monkeypatch.setenv("ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", true_val)
251+
with warnings.catch_warnings(record=True) as w:
252+
warnings.simplefilter("always")
253+
result = experimental_fn()
254+
assert result == "executing"
255+
assert len(w) == 0, f"Bypass failed for env value {true_val}"
256+
257+
258+
def test_experimental_class_bypassed_with_env_var(monkeypatch):
259+
"""Experimental class emits no warning when bypass env var is true."""
260+
true_values = ["true", "True", "TRUE", "1", "yes", "YES", "on", "ON"]
261+
for true_val in true_values:
262+
monkeypatch.setenv("ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", true_val)
263+
with warnings.catch_warnings(record=True) as w:
264+
warnings.simplefilter("always")
265+
exp_class = ExperimentalClass()
266+
result = exp_class.run()
267+
assert result == "running experimental"
268+
assert len(w) == 0, f"Bypass failed for env value {true_val}"
269+
270+
271+
def test_experimental_function_not_bypassed_for_false_env_var(monkeypatch):
272+
"""Experimental function still warns for non-true bypass env var values."""
273+
false_values = ["false", "False", "FALSE", "0", "", "no", "off"]
274+
for false_val in false_values:
275+
monkeypatch.setenv("ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", false_val)
276+
with warnings.catch_warnings(record=True) as w:
277+
warnings.simplefilter("always")
278+
experimental_fn()
279+
assert len(w) == 1
280+
assert "[EXPERIMENTAL] experimental_fn:" in str(w[0].message)
281+
282+
283+
def test_experimental_class_not_bypassed_for_false_env_var(monkeypatch):
284+
"""Experimental class still warns for non-true bypass env var values."""
285+
false_values = ["false", "False", "FALSE", "0", "", "no", "off"]
286+
for false_val in false_values:
287+
monkeypatch.setenv("ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", false_val)
288+
with warnings.catch_warnings(record=True) as w:
289+
warnings.simplefilter("always")
290+
ExperimentalClass()
291+
assert len(w) == 1
292+
assert "[EXPERIMENTAL] ExperimentalClass:" in str(w[0].message)
293+
294+
238295
def test_experimental_class_no_parens_warns():
239296
"""Test that experimental class without parentheses shows default warning."""
240297
with warnings.catch_warnings(record=True) as w:

0 commit comments

Comments
 (0)