Skip to content

Commit a06dca3

Browse files
George Wealecopybara-github
authored andcommitted
feat: Add env var to suppress experimental warnings
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: 794740730
1 parent 52284b1 commit a06dca3

File tree

2 files changed

+110
-2
lines changed

2 files changed

+110
-2
lines changed

src/google/adk/utils/feature_decorator.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,10 +70,18 @@ def new_init(self, *args, **kwargs):
7070
bypass_env_var is not None
7171
and os.environ.get(bypass_env_var, "").lower() == "true"
7272
)
73+
# Suppress experimental warnings if env is set
74+
suppress_experimental = (
75+
label.upper() == "EXPERIMENTAL"
76+
and os.environ.get("ADK_DISABLE_EXPERIMENTAL_WARNING", "").lower()
77+
in ("1", "true", "yes", "on")
78+
)
7379

7480
if should_bypass:
7581
# Bypass completely - no warning, no error
7682
pass
83+
elif suppress_experimental:
84+
pass
7785
elif block_usage:
7886
raise RuntimeError(msg)
7987
else:
@@ -92,10 +100,18 @@ def wrapper(*args, **kwargs):
92100
bypass_env_var is not None
93101
and os.environ.get(bypass_env_var, "").lower() == "true"
94102
)
103+
# Suppress experimental warnings if env is set
104+
suppress_experimental = (
105+
label.upper() == "EXPERIMENTAL"
106+
and os.environ.get("ADK_DISABLE_EXPERIMENTAL_WARNING", "").lower()
107+
in ("1", "true", "yes", "on")
108+
)
95109

96110
if should_bypass:
97111
# Bypass completely - no warning, no error
98112
pass
113+
elif suppress_experimental:
114+
pass
99115
elif block_usage:
100116
raise RuntimeError(msg)
101117
else:

tests/unittests/utils/test_feature_decorator.py

Lines changed: 94 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
1+
import importlib
12
import os
23
import tempfile
34
import warnings
5+
from unittest import mock
46

5-
from google.adk.utils.feature_decorator import experimental
6-
from google.adk.utils.feature_decorator import working_in_progress
7+
from google.adk.utils import feature_decorator
8+
from google.adk.utils.feature_decorator import experimental, working_in_progress
79

810

911
@working_in_progress("in complete feature, don't use yet")
@@ -299,3 +301,93 @@ def test_experimental_function_empty_parens_warns():
299301
assert "This feature is experimental and may change or be removed" in str(
300302
w[0].message
301303
)
304+
305+
306+
def test_experimental_function_warning_suppressed_with_env_var():
307+
"""Test experimental function warning suppressed by env var."""
308+
with mock.patch.dict(os.environ, {"ADK_DISABLE_EXPERIMENTAL_WARNING": "true"}):
309+
importlib.reload(feature_decorator)
310+
from google.adk.utils.feature_decorator import experimental
311+
312+
@experimental
313+
def test_fn():
314+
return "executing"
315+
316+
with warnings.catch_warnings(record=True) as w:
317+
warnings.simplefilter("always")
318+
result = test_fn()
319+
assert result == "executing"
320+
assert len(w) == 0
321+
importlib.reload(feature_decorator) # Reload to clean up
322+
323+
324+
def test_experimental_class_warning_suppressed_with_env_var():
325+
"""Test experimental class warning suppressed by env var."""
326+
with mock.patch.dict(os.environ, {"ADK_DISABLE_EXPERIMENTAL_WARNING": "1"}):
327+
importlib.reload(feature_decorator)
328+
from google.adk.utils.feature_decorator import experimental
329+
330+
@experimental
331+
class TestClass:
332+
def run(self):
333+
return "running experimental"
334+
335+
with warnings.catch_warnings(record=True) as w:
336+
warnings.simplefilter("always")
337+
result = TestClass().run()
338+
assert result == "running experimental"
339+
assert len(w) == 0
340+
importlib.reload(feature_decorator) # Reload to clean up
341+
342+
343+
def test_experimental_env_var_case_insensitive_and_truthy_values():
344+
"""Test experimental warning suppression with various truthy values."""
345+
for val in ["true", "TRUE", "Yes", "on", "1"]:
346+
with mock.patch.dict(os.environ, {"ADK_DISABLE_EXPERIMENTAL_WARNING": val}):
347+
importlib.reload(feature_decorator)
348+
from google.adk.utils.feature_decorator import experimental
349+
350+
@experimental
351+
def test_fn():
352+
return "executing"
353+
354+
with warnings.catch_warnings(record=True) as w:
355+
warnings.simplefilter("always")
356+
_ = test_fn()
357+
assert len(w) == 0
358+
importlib.reload(feature_decorator) # Reload to clean up
359+
360+
361+
def test_experimental_suppression_loads_from_dotenv_file(tmp_path):
362+
"""Test experimental warning suppression loading from .env file."""
363+
try:
364+
from dotenv import load_dotenv
365+
except ImportError:
366+
import pytest
367+
pytest.skip("python-dotenv not available")
368+
369+
env_file = tmp_path / ".env"
370+
env_file.write_text("ADK_DISABLE_EXPERIMENTAL_WARNING=true\n")
371+
372+
# Clear the env var from os.environ if present
373+
with mock.patch.dict(os.environ, {}):
374+
os.environ.pop("ADK_DISABLE_EXPERIMENTAL_WARNING", None)
375+
importlib.reload(feature_decorator)
376+
377+
load_dotenv(env_file.as_posix())
378+
# We need to reload again *after* load_dotenv has modified os.environ
379+
importlib.reload(feature_decorator)
380+
from google.adk.utils.feature_decorator import experimental
381+
382+
@experimental
383+
def test_fn():
384+
return "executing"
385+
386+
try:
387+
with warnings.catch_warnings(record=True) as w:
388+
warnings.simplefilter("always")
389+
_ = test_fn()
390+
assert len(w) == 0
391+
finally:
392+
os.environ.pop("ADK_DISABLE_EXPERIMENTAL_WARNING", None)
393+
importlib.reload(feature_decorator) # Reload to clean up

0 commit comments

Comments
 (0)