diff --git a/src/google/adk/features/__init__.py b/src/google/adk/features/__init__.py index 0a2669d7..f948c077 100644 --- a/src/google/adk/features/__init__.py +++ b/src/google/adk/features/__init__.py @@ -11,3 +11,17 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. + +from ._feature_decorator import experimental +from ._feature_decorator import stable +from ._feature_decorator import working_in_progress +from ._feature_registry import FeatureName +from ._feature_registry import is_feature_enabled + +__all__ = [ + "experimental", + "stable", + "working_in_progress", + "FeatureName", + "is_feature_enabled", +] diff --git a/src/google/adk/features/_feature_decorator.py b/src/google/adk/features/_feature_decorator.py new file mode 100644 index 00000000..dae04758 --- /dev/null +++ b/src/google/adk/features/_feature_decorator.py @@ -0,0 +1,118 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import functools +from typing import Callable +from typing import cast +from typing import TypeVar +from typing import Union + +from ._feature_registry import _get_feature_config +from ._feature_registry import _register_feature +from ._feature_registry import FeatureConfig +from ._feature_registry import FeatureName +from ._feature_registry import FeatureStage +from ._feature_registry import is_feature_enabled + +T = TypeVar("T", bound=Union[Callable, type]) + + +def _make_feature_decorator( + *, + feature_name: FeatureName, + feature_stage: FeatureStage, + default_on: bool = False, +) -> Callable[[T], T]: + """Decorator for experimental features. + + Args: + feature_name: The name of the feature to decorate. + feature_stage: The stage of the feature. + default_on: Whether the feature is enabled by default. + + Returns: + A decorator that checks if the feature is enabled and raises an error if + not. + """ + config = _get_feature_config(feature_name) + if config is None: + config = FeatureConfig(feature_stage, default_on=default_on) + _register_feature(feature_name, config) + + if config.stage != feature_stage: + raise ValueError( + f"Feature '{feature_name}' is being defined with stage" + f" '{feature_stage}', but it was previously registered with stage" + f" '{config.stage}'. Please ensure the feature is consistently defined." + ) + + def decorator(obj: T) -> T: + def check_feature_enabled(): + if not is_feature_enabled(feature_name): + raise RuntimeError(f"Feature {feature_name} is not enabled.") + + if isinstance(obj, type): # decorating a class + original_init = obj.__init__ + + @functools.wraps(original_init) + def new_init(*args, **kwargs): + check_feature_enabled() + return original_init(*args, **kwargs) + + obj.__init__ = new_init + return cast(T, obj) + elif isinstance(obj, Callable): # decorating a function + + @functools.wraps(obj) + def wrapper(*args, **kwargs): + check_feature_enabled() + return obj(*args, **kwargs) + + return cast(T, wrapper) + + else: + raise TypeError( + "@experimental can only be applied to classes or callable objects" + ) + + return decorator + + +def working_in_progress(feature_name: FeatureName) -> Callable[[T], T]: + """Decorator for working in progress features.""" + return _make_feature_decorator( + feature_name=feature_name, + feature_stage=FeatureStage.WIP, + default_on=False, + ) + + +def experimental(feature_name: FeatureName) -> Callable[[T], T]: + """Decorator for experimental features.""" + return _make_feature_decorator( + feature_name=feature_name, + feature_stage=FeatureStage.EXPERIMENTAL, + default_on=False, + ) + + +def stable(feature_name: FeatureName) -> Callable[[T], T]: + """Decorator for stable features.""" + return _make_feature_decorator( + feature_name=feature_name, + feature_stage=FeatureStage.STABLE, + default_on=True, + ) diff --git a/src/google/adk/features/feature_registry.py b/src/google/adk/features/_feature_registry.py similarity index 95% rename from src/google/adk/features/feature_registry.py rename to src/google/adk/features/_feature_registry.py index 0d3740c8..9799d38d 100644 --- a/src/google/adk/features/feature_registry.py +++ b/src/google/adk/features/_feature_registry.py @@ -125,8 +125,13 @@ def is_feature_enabled(feature_name: FeatureName) -> bool: raise ValueError(f"Feature {feature_name} is not registered.") # Check environment variables first (highest priority) - enable_var = f"ADK_ENABLE_{feature_name}" - disable_var = f"ADK_DISABLE_{feature_name}" + feature_name_str = ( + feature_name.value + if isinstance(feature_name, FeatureName) + else feature_name + ) + enable_var = f"ADK_ENABLE_{feature_name_str}" + disable_var = f"ADK_DISABLE_{feature_name_str}" if is_env_enabled(enable_var): if config.stage != FeatureStage.STABLE: _emit_non_stable_warning_once(feature_name, config.stage) diff --git a/tests/unittests/features/test_feature_decorator.py b/tests/unittests/features/test_feature_decorator.py new file mode 100644 index 00000000..54f66e90 --- /dev/null +++ b/tests/unittests/features/test_feature_decorator.py @@ -0,0 +1,207 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import warnings + +from google.adk.features._feature_decorator import experimental +from google.adk.features._feature_decorator import stable +from google.adk.features._feature_decorator import working_in_progress +from google.adk.features._feature_registry import _FEATURE_REGISTRY +from google.adk.features._feature_registry import _get_feature_config +from google.adk.features._feature_registry import _register_feature +from google.adk.features._feature_registry import _WARNED_FEATURES +from google.adk.features._feature_registry import FeatureConfig +from google.adk.features._feature_registry import FeatureStage +import pytest + + +@working_in_progress("WIP_CLASS") +class IncompleteFeature: + + def run(self): + return "running" + + +@working_in_progress("WIP_FUNCTION") +def wip_function(): + return "executing" + + +@experimental("EXPERIMENTAL_CLASS") +class ExperimentalClass: + + def run(self): + return "running" + + +@experimental("EXPERIMENTAL_FUNCTION") +def experimental_function(): + return "executing" + + +@stable("STABLE_CLASS") +class StableClass: + + def run(self): + return "running" + + +@stable("STABLE_FUNCTION") +def stable_function(): + return "executing" + + +@pytest.fixture(autouse=True) +def reset_env_and_registry(monkeypatch): + """Reset environment variables and registry before each test.""" + # Clean up environment variables + for key in list(os.environ.keys()): + if key.startswith("ADK_ENABLE_") or key.startswith("ADK_DISABLE_"): + monkeypatch.delenv(key, raising=False) + + # Add an existing feature to the registry + _register_feature( + "ENABLED_EXPERIMENTAL_FEATURE", + FeatureConfig(FeatureStage.EXPERIMENTAL, default_on=True), + ) + + _register_feature( + "EXPERIMENTAL_FUNCTION", + FeatureConfig(FeatureStage.EXPERIMENTAL, default_on=True), + ) + + +def test_working_in_progress_stage_mismatch(): + """Test that working_in_progress is used with a non-WIP stage.""" + try: + + @working_in_progress("ENABLED_EXPERIMENTAL_FEATURE") + def unused_function(): # pylint: disable=unused-variable + return "unused" + + assert False, "Expected ValueError to be raised." + except ValueError as e: + assert ( + "Feature 'ENABLED_EXPERIMENTAL_FEATURE' is being defined with stage" + " 'FeatureStage.WIP', but it was previously registered with stage" + " 'FeatureStage.EXPERIMENTAL'." + in str(e) + ) + + +def test_working_in_progress_class_raises_error(): + """Test that WIP class raises RuntimeError by default.""" + + try: + IncompleteFeature() + assert False, "Expected RuntimeError to be raised." + except RuntimeError as e: + assert "Feature WIP_CLASS is not enabled." in str(e) + + +def test_working_in_progress_class_bypass_with_env_var(monkeypatch): + """Test that WIP class can be bypassed with env var.""" + + monkeypatch.setenv("ADK_ENABLE_WIP_CLASS", "true") + + with warnings.catch_warnings(record=True) as w: + feature = IncompleteFeature() + feature.run() + assert len(w) == 1 + assert "[WIP] feature WIP_CLASS is enabled." in str(w[0].message) + + +def test_working_in_progress_function_raises_error(): + """Test that WIP function raises RuntimeError by default.""" + + try: + wip_function() + assert False, "Expected RuntimeError to be raised." + except RuntimeError as e: + assert "Feature WIP_FUNCTION is not enabled." in str(e) + + +def test_working_in_progress_function_bypass_with_env_var(monkeypatch): + """Test that WIP function can be bypassed with env var.""" + + monkeypatch.setenv("ADK_ENABLE_WIP_FUNCTION", "true") + + with warnings.catch_warnings(record=True) as w: + wip_function() + assert len(w) == 1 + assert "[WIP] feature WIP_FUNCTION is enabled." in str(w[0].message) + + +def test_disabled_experimental_class_raises_error(): + """Test that disabled experimental class raises RuntimeError by default.""" + + try: + ExperimentalClass() + assert False, "Expected RuntimeError to be raised." + except RuntimeError as e: + assert "Feature EXPERIMENTAL_CLASS is not enabled." in str(e) + + +def test_disabled_experimental_class_bypass_with_env_var(monkeypatch): + """Test that disabled experimental class can be bypassed with env var.""" + + monkeypatch.setenv("ADK_ENABLE_EXPERIMENTAL_CLASS", "true") + + with warnings.catch_warnings(record=True) as w: + feature = ExperimentalClass() + feature.run() + assert len(w) == 1 + assert "[EXPERIMENTAL] feature EXPERIMENTAL_CLASS is enabled." in str( + w[0].message + ) + + +def test_enabled_experimental_function_does_not_raise_error(): + """Test that enabled experimental function does not raise error.""" + + with warnings.catch_warnings(record=True) as w: + experimental_function() + assert len(w) == 1 + assert "[EXPERIMENTAL] feature EXPERIMENTAL_FUNCTION is enabled." in str( + w[0].message + ) + + +def test_enabled_experimental_function_disabled_by_env_var(monkeypatch): + """Test that enabled experimental function can be disabled by env var.""" + + monkeypatch.setenv("ADK_DISABLE_EXPERIMENTAL_FUNCTION", "true") + + try: + experimental_function() + assert False, "Expected RuntimeError to be raised." + except RuntimeError as e: + assert "Feature EXPERIMENTAL_FUNCTION is not enabled." in str(e) + + +def test_stable_class_does_not_raise_error_or_warn(): + """Test that stable class does not raise error or warn.""" + + with warnings.catch_warnings(record=True) as w: + StableClass().run() + assert not w + + +def test_stable_function_does_not_raise_error_or_warn(): + """Test that stable function does not raise error or warn.""" + + with warnings.catch_warnings(record=True) as w: + stable_function() + assert not w diff --git a/tests/unittests/features/test_feature_registry.py b/tests/unittests/features/test_feature_registry.py index 33534172..ab84d986 100644 --- a/tests/unittests/features/test_feature_registry.py +++ b/tests/unittests/features/test_feature_registry.py @@ -17,13 +17,13 @@ from __future__ import annotations import os import warnings -from google.adk.features.feature_registry import _FEATURE_REGISTRY -from google.adk.features.feature_registry import _get_feature_config -from google.adk.features.feature_registry import _register_feature -from google.adk.features.feature_registry import _WARNED_FEATURES -from google.adk.features.feature_registry import FeatureConfig -from google.adk.features.feature_registry import FeatureStage -from google.adk.features.feature_registry import is_feature_enabled +from google.adk.features._feature_registry import _FEATURE_REGISTRY +from google.adk.features._feature_registry import _get_feature_config +from google.adk.features._feature_registry import _register_feature +from google.adk.features._feature_registry import _WARNED_FEATURES +from google.adk.features._feature_registry import FeatureConfig +from google.adk.features._feature_registry import FeatureStage +from google.adk.features._feature_registry import is_feature_enabled import pytest FEATURE_CONFIG_WIP = FeatureConfig(FeatureStage.WIP, default_on=False) @@ -44,17 +44,11 @@ def reset_env_and_registry(monkeypatch): if key.startswith("ADK_ENABLE_") or key.startswith("ADK_DISABLE_"): monkeypatch.delenv(key, raising=False) - # Clear registry (but keep it as a dict for adding test entries) - _FEATURE_REGISTRY.clear() - # Reset warned features set _WARNED_FEATURES.clear() yield - # Clean up after test - _FEATURE_REGISTRY.clear() - # Reset warned features set _WARNED_FEATURES.clear()